use anyhow::Result;
use chainfire_client::Client;
use chrono::Utc;
use clap::Args;
use deployer_types::{
    ClusterNodeRecord, CommissionState, DesiredSystemSpec, HostDeploymentSelector,
    HostDeploymentSpec, HostDeploymentStatus, InstallState, ObservedSystemState,
    ServiceInstanceSpec,
};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::time::Duration;
use tokio::time::sleep;
use tracing::{info, warn};

fn cluster_prefix(cluster_namespace: &str, cluster_id: &str) -> String {
    format!("{}/clusters/{}/", cluster_namespace, cluster_id)
}

fn key_node(cluster_namespace: &str, cluster_id: &str, node_id: &str) -> Vec<u8> {
    format!(
        "{}nodes/{}",
        cluster_prefix(cluster_namespace, cluster_id),
        node_id
    )
    .into_bytes()
}

fn key_desired_system(cluster_namespace: &str, cluster_id: &str, node_id: &str) -> Vec<u8> {
    format!(
        "{}nodes/{}/desired-system",
        cluster_prefix(cluster_namespace, cluster_id),
        node_id
    )
    .into_bytes()
}

fn key_host_deployment_status(
    cluster_namespace: &str,
    cluster_id: &str,
    deployment_name: &str,
) -> Vec<u8> {
    format!(
        "{}deployments/hosts/{}/status",
        cluster_prefix(cluster_namespace, cluster_id),
        deployment_name
    )
    .into_bytes()
}

#[derive(Debug, Clone, Args)]
pub struct HostsCommand {
    #[arg(long)]
    pub endpoint: String,

    #[arg(long, default_value = "photoncloud")]
    pub cluster_namespace: String,

    #[arg(long)]
    pub cluster_id: String,

    #[arg(long, default_value_t = 15)]
    pub interval_secs: u64,

    #[arg(long, default_value_t = 300)]
    pub heartbeat_timeout_secs: u64,

    #[arg(long, default_value_t = false)]
    pub dry_run: bool,

    #[arg(long, default_value_t = false)]
    pub once: bool,
}

pub async fn run(command: HostsCommand) -> Result<()> {
    let controller = HostDeploymentController::new(command);
    if controller.once {
        controller.reconcile_once().await
    } else {
        loop {
            if let Err(error) = controller.reconcile_once().await {
                warn!(error = %error, "host deployment reconciliation failed");
            }
            sleep(controller.interval).await;
        }
    }
}

struct HostDeploymentController {
    endpoint: String,
    cluster_namespace: String,
    cluster_id: String,
    interval: Duration,
    heartbeat_timeout_secs: u64,
    dry_run: bool,
    once: bool,
}

impl HostDeploymentController {
    fn new(command: HostsCommand) -> Self {
        Self {
            endpoint: command.endpoint,
            cluster_namespace: command.cluster_namespace,
            cluster_id: command.cluster_id,
            interval: Duration::from_secs(command.interval_secs),
            heartbeat_timeout_secs: command.heartbeat_timeout_secs,
            dry_run: command.dry_run,
            once: command.once,
        }
    }

    async fn reconcile_once(&self) -> Result<()> {
        let mut client = Client::connect(self.endpoint.clone()).await?;
        let nodes = self.load_nodes(&mut client).await?;
        let desired_systems = self.load_desired_systems(&mut client).await?;
        let observed_systems = self.load_observed_systems(&mut client).await?;
        let instances = self.load_instances(&mut client).await?;
        let deployments = self.load_host_deployments(&mut client).await?;
        let statuses = self.load_host_deployment_statuses(&mut client).await?;

        info!(
            nodes = nodes.len(),
            deployments = deployments.len(),
            instances = instances.len(),
            "loaded host deployment inputs"
        );

        for deployment in deployments {
            let existing_status = statuses.get(&deployment.name).cloned();
            let plan = plan_host_deployment(
                &deployment,
                existing_status.as_ref(),
                &nodes,
                &desired_systems,
                &observed_systems,
                &instances,
                self.heartbeat_timeout_secs,
            );

            if self.dry_run {
                info!(
                    deployment = %deployment.name,
                    phase = plan.status.phase.as_deref().unwrap_or("unknown"),
                    desired_upserts = plan.desired_upserts.len(),
                    desired_deletes = plan.desired_deletes.len(),
                    node_updates = plan.node_updates.len(),
                    "would reconcile host deployment"
                );
                continue;
            }

            for desired in &plan.desired_upserts {
                client
                    .put(
                        &key_desired_system(
                            &self.cluster_namespace,
                            &self.cluster_id,
                            &desired.node_id,
                        ),
                        &serde_json::to_vec(desired)?,
                    )
                    .await?;
            }

            for node_id in &plan.desired_deletes {
                client
                    .delete(&key_desired_system(
                        &self.cluster_namespace,
                        &self.cluster_id,
                        node_id,
                    ))
                    .await?;
            }

            for node in plan.node_updates.values() {
                client
                    .put(
                        &key_node(&self.cluster_namespace, &self.cluster_id, &node.node_id),
                        &serde_json::to_vec(node)?,
                    )
                    .await?;
            }

            client
                .put(
                    &key_host_deployment_status(
                        &self.cluster_namespace,
                        &self.cluster_id,
                        &deployment.name,
                    ),
                    &serde_json::to_vec(&plan.status)?,
                )
                .await?;
        }

        Ok(())
    }

    async fn load_nodes(&self, client: &mut Client) -> Result<Vec<ClusterNodeRecord>> {
        let prefix = format!(
            "{}nodes/",
            cluster_prefix(&self.cluster_namespace, &self.cluster_id)
        );
        let kvs = client.get_prefix(prefix.as_bytes()).await?;
        let mut nodes = Vec::new();

        for (key, value) in kvs {
            let key = String::from_utf8_lossy(&key);
            let Some(suffix) = key.strip_prefix(&prefix) else {
                continue;
            };
            if suffix.contains('/') {
                continue;
            }
            match serde_json::from_slice::<ClusterNodeRecord>(&value) {
                Ok(node) => nodes.push(node),
                Err(error) => warn!(error = %error, key = %key, "failed to decode cluster node"),
            }
        }

        nodes.sort_by(|lhs, rhs| lhs.node_id.cmp(&rhs.node_id));
        Ok(nodes)
    }

    async fn load_desired_systems(
        &self,
        client: &mut Client,
    ) -> Result<HashMap<String, DesiredSystemSpec>> {
        let prefix = format!(
            "{}nodes/",
            cluster_prefix(&self.cluster_namespace, &self.cluster_id)
        );
        let kvs = client.get_prefix(prefix.as_bytes()).await?;
        let mut desired = HashMap::new();

        for (key, value) in kvs {
            let key = String::from_utf8_lossy(&key);
            if !key.ends_with("/desired-system") {
                continue;
            }
            match serde_json::from_slice::<DesiredSystemSpec>(&value) {
                Ok(spec) => {
                    desired.insert(spec.node_id.clone(), spec);
                }
                Err(error) => warn!(error = %error, key = %key, "failed to decode desired-system"),
            }
        }

        Ok(desired)
    }

    async fn load_observed_systems(
        &self,
        client: &mut Client,
    ) -> Result<HashMap<String, ObservedSystemState>> {
        let prefix = format!(
            "{}nodes/",
            cluster_prefix(&self.cluster_namespace, &self.cluster_id)
        );
        let kvs = client.get_prefix(prefix.as_bytes()).await?;
        let mut observed = HashMap::new();

        for (key, value) in kvs {
            let key = String::from_utf8_lossy(&key);
            if !key.ends_with("/observed-system") {
                continue;
            }
            match serde_json::from_slice::<ObservedSystemState>(&value) {
                Ok(state) => {
                    observed.insert(state.node_id.clone(), state);
                }
                Err(error) => warn!(error = %error, key = %key, "failed to decode observed-system"),
            }
        }

        Ok(observed)
    }

    async fn load_instances(&self, client: &mut Client) -> Result<Vec<ServiceInstanceSpec>> {
        let prefix = format!(
            "{}instances/",
            cluster_prefix(&self.cluster_namespace, &self.cluster_id)
        );
        let kvs = client.get_prefix(prefix.as_bytes()).await?;
        let mut instances = Vec::new();

        for (key, value) in kvs {
            let key = String::from_utf8_lossy(&key);
            match serde_json::from_slice::<ServiceInstanceSpec>(&value) {
                Ok(instance) => instances.push(instance),
                Err(error) => {
                    warn!(error = %error, key = %key, "failed to decode service instance")
                }
            }
        }

        Ok(instances)
    }

    async fn load_host_deployments(&self, client: &mut Client) -> Result<Vec<HostDeploymentSpec>> {
        let prefix = format!(
            "{}deployments/hosts/",
            cluster_prefix(&self.cluster_namespace, &self.cluster_id)
        );
        let kvs = client.get_prefix(prefix.as_bytes()).await?;
        let mut deployments = Vec::new();

        for (key, value) in kvs {
            let key = String::from_utf8_lossy(&key);
            if !key.ends_with("/spec") {
                continue;
            }
            match serde_json::from_slice::<HostDeploymentSpec>(&value) {
                Ok(spec) => deployments.push(spec),
                Err(error) => warn!(error = %error, key = %key, "failed to decode host deployment"),
            }
        }

        deployments.sort_by(|lhs, rhs| lhs.name.cmp(&rhs.name));
        Ok(deployments)
    }

    async fn load_host_deployment_statuses(
        &self,
        client: &mut Client,
    ) -> Result<HashMap<String, HostDeploymentStatus>> {
        let prefix = format!(
            "{}deployments/hosts/",
            cluster_prefix(&self.cluster_namespace, &self.cluster_id)
        );
        let kvs = client.get_prefix(prefix.as_bytes()).await?;
        let mut statuses = HashMap::new();

        for (key, value) in kvs {
            let key = String::from_utf8_lossy(&key);
            if !key.ends_with("/status") {
                continue;
            }
            match serde_json::from_slice::<HostDeploymentStatus>(&value) {
                Ok(status) => {
                    statuses.insert(status.name.clone(), status);
                }
                Err(error) => {
                    warn!(error = %error, key = %key, "failed to decode host deployment status")
                }
            }
        }

        Ok(statuses)
    }
}

#[derive(Debug, Default)]
struct HostDeploymentPlan {
    status: HostDeploymentStatus,
    desired_upserts: Vec<DesiredSystemSpec>,
    desired_deletes: Vec<String>,
    node_updates: BTreeMap<String, ClusterNodeRecord>,
}

fn plan_host_deployment(
    deployment: &HostDeploymentSpec,
    existing_status: Option<&HostDeploymentStatus>,
    nodes: &[ClusterNodeRecord],
    desired_systems: &HashMap<String, DesiredSystemSpec>,
    observed_systems: &HashMap<String, ObservedSystemState>,
    instances: &[ServiceInstanceSpec],
    heartbeat_timeout_secs: u64,
) -> HostDeploymentPlan {
    let now = Utc::now();
    let target_configuration = deployment.nixos_configuration.clone();
    let selector_matches = select_nodes(nodes, &deployment.selector);
    let selected_node_ids = selector_matches
        .iter()
        .map(|node| node.node_id.clone())
        .collect::<HashSet<_>>();
    let instance_counts = active_instances_per_node(instances);
    let mut completed = Vec::new();
    let mut in_progress = Vec::new();
    let mut failed = Vec::new();
    let mut eligible_candidates = Vec::new();
    let mut desired_upserts = Vec::new();
    let mut node_updates = BTreeMap::new();
    let batch_size = deployment.batch_size.unwrap_or(1).max(1) as usize;
    let max_unavailable = deployment.max_unavailable.unwrap_or(1).max(1) as usize;
    let operator_paused = existing_status
        .map(|status| status.paused_by_operator)
        .unwrap_or(false);
    let spec_paused = deployment.paused.unwrap_or(false);
    let mut desired_deletes = desired_systems
        .iter()
        .filter(|(node_id, desired)| {
            desired.deployment_id.as_deref() == Some(deployment.name.as_str())
                && !selected_node_ids.contains(node_id.as_str())
        })
        .map(|(node_id, _)| node_id.clone())
        .collect::<Vec<_>>();

    for node in &selector_matches {
        let desired = desired_systems.get(&node.node_id);
        let observed = observed_systems.get(&node.node_id);
        let is_completed = is_node_completed(
            deployment,
            node,
            desired,
            observed,
            target_configuration.as_deref(),
        );
        let is_failed = is_node_failed(deployment, desired, observed);
        let is_in_progress =
            is_node_in_progress(deployment, desired, observed, is_completed, is_failed)
                || (deployment.drain_before_apply == Some(true)
                    && node.state.as_deref() == Some("draining")
                    && instance_counts
                        .get(&node.node_id)
                        .copied()
                        .unwrap_or_default()
                        > 0);

        if is_completed {
            completed.push(node.node_id.clone());
            if deployment.drain_before_apply == Some(true)
                && node.state.as_deref() == Some("draining")
            {
                let mut updated = (*node).clone();
                updated.state = Some("active".to_string());
                node_updates.insert(updated.node_id.clone(), updated);
            }
            continue;
        }

        if is_failed {
            failed.push(node.node_id.clone());
            continue;
        }

        if is_in_progress {
            in_progress.push(node.node_id.clone());
            continue;
        }

        if node_is_rollout_candidate(node, heartbeat_timeout_secs) {
            eligible_candidates.push((*node).clone());
        }
    }

    let unavailable = in_progress.len() + failed.len();
    let paused = operator_paused || spec_paused || !failed.is_empty();
    let remaining_unavailable_budget = max_unavailable.saturating_sub(unavailable);
    let remaining_batch_budget = batch_size.saturating_sub(in_progress.len());
    let max_starts =
        if deployment.nixos_configuration.is_some() || deployment.target_system.is_some() {
            remaining_unavailable_budget.min(remaining_batch_budget)
        } else {
            0
        };
    let mut planned = 0usize;
    let mut newly_started = Vec::new();

    if !paused && max_starts > 0 {
        for node in eligible_candidates {
            if planned >= max_starts {
                break;
            }

            let remaining_instances = instance_counts
                .get(&node.node_id)
                .copied()
                .unwrap_or_default();
            if deployment.drain_before_apply == Some(true) && remaining_instances > 0 {
                let mut updated = node.clone();
                updated.state = Some("draining".to_string());
                node_updates.insert(updated.node_id.clone(), updated);
                in_progress.push(node.node_id.clone());
                newly_started.push(node.node_id.clone());
                planned += 1;
                continue;
            }

            let desired = DesiredSystemSpec {
                node_id: node.node_id.clone(),
                deployment_id: Some(deployment.name.clone()),
                nixos_configuration: deployment.nixos_configuration.clone(),
                target_system: deployment.target_system.clone(),
                flake_ref: deployment.flake_ref.clone(),
                switch_action: deployment
                    .switch_action
                    .clone()
                    .or_else(|| Some("switch".to_string())),
                health_check_command: deployment.health_check_command.clone(),
                rollback_on_failure: Some(deployment.rollback_on_failure.unwrap_or(true)),
                drain_before_apply: Some(deployment.drain_before_apply.unwrap_or(false)),
            };
            newly_started.push(node.node_id.clone());
            in_progress.push(node.node_id.clone());
            planned += 1;
            if deployment.drain_before_apply == Some(true)
                && node.state.as_deref() != Some("draining")
            {
                let mut updated = node.clone();
                updated.state = Some("draining".to_string());
                node_updates.insert(updated.node_id.clone(), updated);
            }
            desired_upserts.push(desired);
        }
    }

    let mut status = existing_status.cloned().unwrap_or_default();
    status.name = deployment.name.clone();
    status.selected_nodes = selector_matches
        .iter()
        .map(|node| node.node_id.clone())
        .collect();
    status.completed_nodes = dedup_sorted(completed);
    status.in_progress_nodes = dedup_sorted(in_progress);
    status.failed_nodes = dedup_sorted(failed);
    status.paused_by_operator = operator_paused;
    status.paused = paused;
    status.phase = Some(
        if status.selected_nodes.is_empty() {
            "idle"
        } else if deployment.nixos_configuration.is_none() && deployment.target_system.is_none() {
            "invalid"
        } else if status.paused {
            "paused"
        } else if status.completed_nodes.len() == status.selected_nodes.len() {
            "completed"
        } else if !newly_started.is_empty() || !status.in_progress_nodes.is_empty() {
            "running"
        } else {
            "ready"
        }
        .to_string(),
    );
    status.message = Some(format!(
        "selected={} completed={} in_progress={} failed={} newly_started={}",
        status.selected_nodes.len(),
        status.completed_nodes.len(),
        status.in_progress_nodes.len(),
        status.failed_nodes.len(),
        newly_started.len()
    ));
    status.updated_at = Some(now);

    HostDeploymentPlan {
        status,
        desired_upserts,
        desired_deletes: {
            desired_deletes.sort();
            desired_deletes.dedup();
            desired_deletes
        },
        node_updates,
    }
}

fn select_nodes<'a>(
    nodes: &'a [ClusterNodeRecord],
    selector: &HostDeploymentSelector,
) -> Vec<&'a ClusterNodeRecord> {
    let explicit_nodes = selector.node_ids.iter().collect::<HashSet<_>>();
    let explicit_mode = !explicit_nodes.is_empty();
    let mut selected = nodes
        .iter()
        .filter(|node| {
            (!explicit_mode || explicit_nodes.contains(&node.node_id))
                && (selector.roles.is_empty()
                    || node
                        .roles
                        .iter()
                        .any(|role| selector.roles.iter().any(|expected| expected == role)))
                && (selector.pools.is_empty()
                    || node
                        .pool
                        .as_deref()
                        .map(|pool| selector.pools.iter().any(|expected| expected == pool))
                        .unwrap_or(false))
                && (selector.node_classes.is_empty()
                    || node
                        .node_class
                        .as_deref()
                        .map(|node_class| {
                            selector
                                .node_classes
                                .iter()
                                .any(|expected| expected == node_class)
                        })
                        .unwrap_or(false))
                && selector
                    .match_labels
                    .iter()
                    .all(|(key, value)| node.labels.get(key) == Some(value))
        })
        .collect::<Vec<_>>();
    selected.sort_by(|lhs, rhs| lhs.node_id.cmp(&rhs.node_id));
    selected
}

fn active_instances_per_node(instances: &[ServiceInstanceSpec]) -> HashMap<String, usize> {
    let mut counts = HashMap::new();
    for instance in instances {
        if matches!(instance.state.as_deref(), Some("failed") | Some("deleted")) {
            continue;
        }
        *counts.entry(instance.node_id.clone()).or_insert(0usize) += 1;
    }
    counts
}

fn node_is_rollout_candidate(node: &ClusterNodeRecord, heartbeat_timeout_secs: u64) -> bool {
    if matches!(
        node.commission_state,
        Some(CommissionState::Discovered | CommissionState::Commissioning)
    ) {
        return false;
    }
    if matches!(
        node.install_state,
        Some(InstallState::Installing | InstallState::Failed | InstallState::ReinstallRequested)
    ) {
        return false;
    }
    if !matches!(node.state.as_deref(), Some("active") | Some("draining")) {
        return false;
    }
    if heartbeat_timeout_secs == 0 {
        return true;
    }
    let Some(last) = node.last_heartbeat else {
        return true;
    };
    Utc::now().signed_duration_since(last).num_seconds() <= heartbeat_timeout_secs as i64
}

fn is_node_completed(
    deployment: &HostDeploymentSpec,
    _node: &ClusterNodeRecord,
    desired: Option<&DesiredSystemSpec>,
    observed: Option<&ObservedSystemState>,
    target_configuration: Option<&str>,
) -> bool {
    observed
        .filter(|observed| observed.status.as_deref() == Some("active"))
        .map(|observed| {
            target_configuration
                .map(|target| observed.nixos_configuration.as_deref() == Some(target))
                .or_else(|| {
                    deployment
                        .target_system
                        .as_deref()
                        .map(|target| observed.target_system.as_deref() == Some(target))
                })
                .unwrap_or(false)
        })
        .unwrap_or(false)
        && desired
            .and_then(|desired| desired.deployment_id.as_deref())
            .map(|deployment_id| deployment_id == deployment.name)
            .unwrap_or(false)
}

fn is_node_failed(
    deployment: &HostDeploymentSpec,
    desired: Option<&DesiredSystemSpec>,
    observed: Option<&ObservedSystemState>,
) -> bool {
    desired
        .and_then(|desired| desired.deployment_id.as_deref())
        .map(|deployment_id| deployment_id == deployment.name)
        .unwrap_or(false)
        && observed
            .and_then(|observed| observed.status.as_deref())
            .map(|status| matches!(status, "failed" | "rolled-back"))
            .unwrap_or(false)
}

fn is_node_in_progress(
    deployment: &HostDeploymentSpec,
    desired: Option<&DesiredSystemSpec>,
    observed: Option<&ObservedSystemState>,
    is_completed: bool,
    is_failed: bool,
) -> bool {
    if is_completed || is_failed {
        return false;
    }
    desired
        .and_then(|desired| desired.deployment_id.as_deref())
        .map(|deployment_id| deployment_id == deployment.name)
        .unwrap_or(false)
        || observed
            .and_then(|observed| observed.status.as_deref())
            .map(|status| {
                matches!(
                    status,
                    "planning" | "pending" | "reconciling" | "verifying" | "staged"
                )
            })
            .unwrap_or(false)
}

fn dedup_sorted(mut values: Vec<String>) -> Vec<String> {
    values.sort();
    values.dedup();
    values
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_node(node_id: &str, failure_domain: &str) -> ClusterNodeRecord {
        ClusterNodeRecord {
            node_id: node_id.to_string(),
            machine_id: None,
            ip: "10.0.0.1".to_string(),
            hostname: node_id.to_string(),
            roles: vec!["worker".to_string()],
            labels: HashMap::from([
                ("tier".to_string(), "general".to_string()),
                ("failure_domain".to_string(), failure_domain.to_string()),
            ]),
            pool: Some("general".to_string()),
            node_class: Some("worker-linux".to_string()),
            failure_domain: Some(failure_domain.to_string()),
            nix_profile: None,
            install_plan: None,
            hardware_facts: None,
            state: Some("active".to_string()),
            commission_state: Some(CommissionState::Commissioned),
            install_state: Some(InstallState::Installed),
            commissioned_at: None,
            last_inventory_hash: None,
            power_state: None,
            bmc_ref: None,
            last_heartbeat: Some(Utc::now()),
        }
    }

    fn test_deployment() -> HostDeploymentSpec {
        HostDeploymentSpec {
            name: "worker-rollout".to_string(),
            selector: HostDeploymentSelector {
                node_ids: vec![],
                roles: vec!["worker".to_string()],
                pools: vec!["general".to_string()],
                node_classes: vec!["worker-linux".to_string()],
                match_labels: HashMap::from([("tier".to_string(), "general".to_string())]),
            },
            nixos_configuration: Some("worker-golden".to_string()),
            target_system: Some("/nix/store/worker-golden".to_string()),
            flake_ref: Some("/opt/plasmacloud-src".to_string()),
            batch_size: Some(1),
            max_unavailable: Some(1),
            health_check_command: vec!["true".to_string()],
            switch_action: Some("switch".to_string()),
            rollback_on_failure: Some(true),
            drain_before_apply: Some(false),
            reboot_policy: None,
            paused: Some(false),
        }
    }

    #[test]
    fn plan_rollout_starts_one_node_per_batch() {
        let deployment = test_deployment();
        let nodes = vec![test_node("node01", "rack-a"), test_node("node02", "rack-b")];
        let plan = plan_host_deployment(
            &deployment,
            None,
            &nodes,
            &HashMap::new(),
            &HashMap::new(),
            &[],
            300,
        );

        assert_eq!(plan.desired_upserts.len(), 1);
        assert_eq!(
            plan.desired_upserts[0].target_system.as_deref(),
            Some("/nix/store/worker-golden")
        );
        assert_eq!(plan.status.in_progress_nodes, vec!["node01".to_string()]);
        assert_eq!(plan.status.phase.as_deref(), Some("running"));
    }

    #[test]
    fn plan_rollout_pauses_on_failed_node() {
        let deployment = test_deployment();
        let nodes = vec![test_node("node01", "rack-a"), test_node("node02", "rack-b")];
        let desired = HashMap::from([(
            "node01".to_string(),
            DesiredSystemSpec {
                node_id: "node01".to_string(),
                deployment_id: Some("worker-rollout".to_string()),
                nixos_configuration: Some("worker-golden".to_string()),
                target_system: Some("/nix/store/worker-golden".to_string()),
                flake_ref: None,
                switch_action: Some("switch".to_string()),
                health_check_command: Vec::new(),
                rollback_on_failure: Some(true),
                drain_before_apply: Some(false),
            },
        )]);
        let observed = HashMap::from([(
            "node01".to_string(),
            ObservedSystemState {
                node_id: "node01".to_string(),
                nixos_configuration: Some("worker-golden".to_string()),
                status: Some("rolled-back".to_string()),
                ..ObservedSystemState::default()
            },
        )]);

        let plan = plan_host_deployment(&deployment, None, &nodes, &desired, &observed, &[], 300);

        assert!(plan.desired_upserts.is_empty());
        assert!(plan.status.paused);
        assert_eq!(plan.status.failed_nodes, vec!["node01".to_string()]);
    }

    #[test]
    fn plan_rollout_drains_before_apply_when_instances_exist() {
        let mut deployment = test_deployment();
        deployment.drain_before_apply = Some(true);
        let nodes = vec![test_node("node01", "rack-a")];
        let instances = vec![ServiceInstanceSpec {
            instance_id: "api-node01".to_string(),
            service: "api".to_string(),
            node_id: "node01".to_string(),
            ip: "10.0.0.1".to_string(),
            port: 8080,
            mesh_port: None,
            version: None,
            health_check: None,
            process: None,
            container: None,
            managed_by: Some("fleet-scheduler".to_string()),
            state: Some("active".to_string()),
            last_heartbeat: None,
            observed_at: None,
        }];

        let plan = plan_host_deployment(
            &deployment,
            None,
            &nodes,
            &HashMap::new(),
            &HashMap::new(),
            &instances,
            300,
        );

        assert!(plan.desired_upserts.is_empty());
        assert_eq!(
            plan.node_updates
                .get("node01")
                .and_then(|node| node.state.as_deref()),
            Some("draining")
        );
        assert_eq!(plan.status.in_progress_nodes, vec!["node01".to_string()]);
    }
}
