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 serde_json::Value;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::time::Duration;
use tokio::sync::{mpsc, oneshot};
use tokio::time::MissedTickBehavior;
use tracing::{info, warn};

use crate::watcher::{ChainfireWatcher, WatchChange};

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

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

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

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

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 = "ultracloud")]
    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 mut controller = HostDeploymentController::new(command);
    controller.run_loop().await
}

struct HostDeploymentController {
    endpoint: String,
    cluster_namespace: String,
    cluster_id: String,
    interval: Duration,
    heartbeat_timeout_secs: u64,
    dry_run: bool,
    once: bool,
    watched_node_signatures: HashMap<String, Vec<u8>>,
    watched_desired_signatures: HashMap<String, Vec<u8>>,
    watched_observed_signatures: HashMap<String, Vec<u8>>,
    watched_instance_signatures: HashMap<String, Vec<u8>>,
    watched_deployment_spec_signatures: HashMap<String, Vec<u8>>,
    watched_deployment_status_signatures: HashMap<String, Vec<u8>>,
    pending_kv_writes: HashMap<String, Vec<u8>>,
    pending_kv_deletes: HashSet<String>,
}

async fn spawn_watch_task(
    watcher: ChainfireWatcher,
    tx: mpsc::Sender<WatchChange>,
    ready: oneshot::Sender<()>,
) {
    let mut ready = Some(ready);
    if let Err(error) = watcher
        .watch_with_ready(
            move || {
                if let Some(ready) = ready.take() {
                    let _ = ready.send(());
                }
                Ok(())
            },
            move |change| {
                tx.try_send(change).map_err(|send_error| {
                    anyhow::anyhow!("failed to enqueue watch change: {send_error}")
                })
            },
        )
        .await
    {
        warn!(error = %error, "host deployment watch task exited");
    }
}

async fn wait_for_watchers_ready(receivers: Vec<oneshot::Receiver<()>>) {
    for receiver in receivers {
        match tokio::time::timeout(Duration::from_secs(10), receiver).await {
            Ok(Ok(())) => {}
            Ok(Err(_)) => warn!("host deployment watch task ended before signaling readiness"),
            Err(_) => warn!("timed out waiting for host deployment watch readiness"),
        }
    }
}

fn canonical_json_signature(bytes: &[u8], drop_keys: &[&str]) -> Option<Vec<u8>> {
    let mut value: Value = serde_json::from_slice(bytes).ok()?;
    strip_observed_fields(&mut value, drop_keys);
    serde_json::to_vec(&sort_json_value(value)).ok()
}

fn strip_observed_fields(value: &mut Value, drop_keys: &[&str]) {
    match value {
        Value::Object(map) => {
            for key in drop_keys {
                map.remove(*key);
            }
            for nested in map.values_mut() {
                strip_observed_fields(nested, drop_keys);
            }
        }
        Value::Array(items) => {
            for item in items {
                strip_observed_fields(item, drop_keys);
            }
        }
        _ => {}
    }
}

fn sort_json_value(value: Value) -> Value {
    match value {
        Value::Array(items) => Value::Array(items.into_iter().map(sort_json_value).collect()),
        Value::Object(map) => {
            let mut keys = map.keys().cloned().collect::<Vec<_>>();
            keys.sort();
            let mut normalized = serde_json::Map::new();
            for key in keys {
                if let Some(value) = map.get(&key) {
                    normalized.insert(key, sort_json_value(value.clone()));
                }
            }
            Value::Object(normalized)
        }
        other => other,
    }
}

fn host_deployment_status_signature(bytes: &[u8]) -> Option<Vec<u8>> {
    let value: Value = serde_json::from_slice(bytes).ok()?;
    serde_json::to_vec(&serde_json::json!({
        "paused_by_operator": value.get("paused_by_operator").and_then(Value::as_bool).unwrap_or(false),
        "aborted": value.get("phase").and_then(Value::as_str) == Some("aborted"),
    }))
    .ok()
}

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,
            watched_node_signatures: HashMap::new(),
            watched_desired_signatures: HashMap::new(),
            watched_observed_signatures: HashMap::new(),
            watched_instance_signatures: HashMap::new(),
            watched_deployment_spec_signatures: HashMap::new(),
            watched_deployment_status_signatures: HashMap::new(),
            pending_kv_writes: HashMap::new(),
            pending_kv_deletes: HashSet::new(),
        }
    }

    async fn run_loop(&mut self) -> Result<()> {
        if self.once {
            return self.reconcile_once().await;
        }

        self.reconcile_once().await?;

        let (watch_tx, mut watch_rx) = mpsc::channel::<WatchChange>(128);
        let ready_receivers = self.spawn_watchers(watch_tx);
        wait_for_watchers_ready(ready_receivers).await;
        self.reconcile_once().await?;
        let mut interval = tokio::time::interval(self.interval);
        interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
        interval.tick().await;

        loop {
            tokio::select! {
                _ = interval.tick() => {
                    if let Err(error) = self.reconcile_once().await {
                        warn!(error = %error, "host deployment reconciliation failed");
                    }
                }
                maybe_change = watch_rx.recv() => {
                    let Some(change) = maybe_change else {
                        continue;
                    };

                    let mut should_reconcile = self.apply_watch_change(&change);
                    while let Ok(extra) = watch_rx.try_recv() {
                        should_reconcile |= self.apply_watch_change(&extra);
                    }

                    if should_reconcile {
                        if let Err(error) = self.reconcile_once().await {
                            warn!(error = %error, "host deployment watch-triggered reconciliation failed");
                        }
                    }
                }
            }
        }
    }

    async fn reconcile_once(&mut 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 {
                let key = key_desired_system(&self.cluster_namespace, &self.cluster_id, &desired.node_id);
                let payload = serde_json::to_vec(desired)?;
                let key_str = self.record_pending_write(&key, payload.clone());
                if let Err(error) = client.put(&key, &payload).await {
                    self.pending_kv_writes.remove(&key_str);
                    return Err(error.into());
                }
            }

            for node_id in &plan.desired_deletes {
                let key = key_desired_system(&self.cluster_namespace, &self.cluster_id, node_id);
                let key_str = self.record_pending_delete(&key);
                if !client.delete(&key).await? {
                    self.pending_kv_deletes.remove(&key_str);
                }
            }

            for node in plan.node_updates.values() {
                let key = key_node(&self.cluster_namespace, &self.cluster_id, &node.node_id);
                let payload = serde_json::to_vec(node)?;
                let key_str = self.record_pending_write(&key, payload.clone());
                if let Err(error) = client.put(&key, &payload).await {
                    self.pending_kv_writes.remove(&key_str);
                    return Err(error.into());
                }
            }

            let key = key_host_deployment_status(
                &self.cluster_namespace,
                &self.cluster_id,
                &deployment.name,
            );
            let payload = serde_json::to_vec(&plan.status)?;
            let key_str = self.record_pending_write(&key, payload.clone());
            if let Err(error) = client.put(&key, &payload).await {
                self.pending_kv_writes.remove(&key_str);
                return Err(error.into());
            }
        }

        Ok(())
    }

    fn spawn_watchers(&self, tx: mpsc::Sender<WatchChange>) -> Vec<oneshot::Receiver<()>> {
        let mut ready_receivers = Vec::new();

        let node_watcher = ChainfireWatcher::prefix(
            self.endpoint.clone(),
            nodes_prefix(&self.cluster_namespace, &self.cluster_id),
        );
        let (node_ready_tx, node_ready_rx) = oneshot::channel();
        tokio::spawn(spawn_watch_task(node_watcher, tx.clone(), node_ready_tx));
        ready_receivers.push(node_ready_rx);

        let instance_watcher = ChainfireWatcher::prefix(
            self.endpoint.clone(),
            instances_prefix(&self.cluster_namespace, &self.cluster_id),
        );
        let (instance_ready_tx, instance_ready_rx) = oneshot::channel();
        tokio::spawn(spawn_watch_task(
            instance_watcher,
            tx.clone(),
            instance_ready_tx,
        ));
        ready_receivers.push(instance_ready_rx);

        let deployment_watcher = ChainfireWatcher::prefix(
            self.endpoint.clone(),
            host_deployments_prefix(&self.cluster_namespace, &self.cluster_id),
        );
        let (deployment_ready_tx, deployment_ready_rx) = oneshot::channel();
        tokio::spawn(spawn_watch_task(deployment_watcher, tx, deployment_ready_tx));
        ready_receivers.push(deployment_ready_rx);

        ready_receivers
    }

    fn record_pending_write(&mut self, key: &[u8], payload: Vec<u8>) -> String {
        let key_str = String::from_utf8_lossy(key).to_string();
        self.pending_kv_deletes.remove(&key_str);
        self.pending_kv_writes.insert(key_str.clone(), payload);
        key_str
    }

    fn record_pending_delete(&mut self, key: &[u8]) -> String {
        let key_str = String::from_utf8_lossy(key).to_string();
        self.pending_kv_writes.remove(&key_str);
        self.pending_kv_deletes.insert(key_str.clone());
        key_str
    }

    fn consume_pending_write(&mut self, key: &str, change: &WatchChange) -> Option<Vec<u8>> {
        let expected = self.pending_kv_writes.remove(key)?;
        if matches!(change.event_type, chainfire_client::EventType::Put) && change.value == expected {
            self.pending_kv_deletes.remove(key);
            return Some(expected);
        }
        None
    }

    fn consume_pending_delete(&mut self, key: &str, change: &WatchChange) -> bool {
        self.pending_kv_deletes.remove(key)
            && matches!(change.event_type, chainfire_client::EventType::Delete)
    }

    fn apply_watch_change(&mut self, change: &WatchChange) -> bool {
        let key = String::from_utf8_lossy(&change.key).to_string();
        let nodes_prefix = String::from_utf8_lossy(&nodes_prefix(&self.cluster_namespace, &self.cluster_id))
            .to_string();
        let instances_prefix =
            String::from_utf8_lossy(&instances_prefix(&self.cluster_namespace, &self.cluster_id))
                .to_string();
        let deployments_prefix = String::from_utf8_lossy(&host_deployments_prefix(
            &self.cluster_namespace,
            &self.cluster_id,
        ))
        .to_string();

        if let Some(suffix) = key.strip_prefix(&nodes_prefix) {
            if suffix.ends_with("/desired-system") {
                return self.apply_desired_system_watch_change(&key, change);
            }
            if suffix.ends_with("/observed-system") {
                return self.apply_observed_system_watch_change(&key, change);
            }
            if !suffix.is_empty() && !suffix.contains('/') {
                return self.apply_node_record_watch_change(&key, change);
            }
            return false;
        }

        if key.starts_with(&instances_prefix) {
            return self.apply_instance_watch_change(&key, change);
        }

        if let Some(suffix) = key.strip_prefix(&deployments_prefix) {
            if suffix.ends_with("/spec") {
                return self.apply_deployment_spec_watch_change(&key, change);
            }
            if suffix.ends_with("/status") {
                return self.apply_deployment_status_watch_change(&key, change);
            }
        }

        false
    }

    fn apply_node_record_watch_change(&mut self, key: &str, change: &WatchChange) -> bool {
        if self.consume_pending_delete(key, change) {
            self.watched_node_signatures.remove(key);
            return false;
        }
        if let Some(payload) = self.consume_pending_write(key, change) {
            if let Some(signature) = canonical_json_signature(&payload, &["last_heartbeat"]) {
                self.watched_node_signatures
                    .insert(key.to_string(), signature);
            } else {
                self.watched_node_signatures.remove(key);
            }
            return false;
        }
        if matches!(change.event_type, chainfire_client::EventType::Delete) {
            self.watched_node_signatures.remove(key);
            return true;
        }

        let Some(signature) = canonical_json_signature(&change.value, &["last_heartbeat"]) else {
            return true;
        };
        if self.watched_node_signatures.get(key) == Some(&signature) {
            return false;
        }
        self.watched_node_signatures
            .insert(key.to_string(), signature);
        true
    }

    fn apply_desired_system_watch_change(&mut self, key: &str, change: &WatchChange) -> bool {
        if self.consume_pending_delete(key, change) {
            self.watched_desired_signatures.remove(key);
            return false;
        }
        if let Some(payload) = self.consume_pending_write(key, change) {
            if let Some(signature) = canonical_json_signature(&payload, &[]) {
                self.watched_desired_signatures
                    .insert(key.to_string(), signature);
            } else {
                self.watched_desired_signatures.remove(key);
            }
            return false;
        }
        if matches!(change.event_type, chainfire_client::EventType::Delete) {
            self.watched_desired_signatures.remove(key);
            return true;
        }

        let Some(signature) = canonical_json_signature(&change.value, &[]) else {
            return true;
        };
        if self.watched_desired_signatures.get(key) == Some(&signature) {
            return false;
        }
        self.watched_desired_signatures
            .insert(key.to_string(), signature);
        true
    }

    fn apply_observed_system_watch_change(&mut self, key: &str, change: &WatchChange) -> bool {
        if matches!(change.event_type, chainfire_client::EventType::Delete) {
            self.watched_observed_signatures.remove(key);
            return true;
        }

        let Some(signature) = canonical_json_signature(
            &change.value,
            &[
                "flake_root",
                "configured_system",
                "current_system",
                "booted_system",
                "rollback_system",
                "switch_action",
                "reboot_required",
                "last_attempt",
                "last_success",
                "last_error",
            ],
        ) else {
            return true;
        };
        if self.watched_observed_signatures.get(key) == Some(&signature) {
            return false;
        }
        self.watched_observed_signatures
            .insert(key.to_string(), signature);
        true
    }

    fn apply_instance_watch_change(&mut self, key: &str, change: &WatchChange) -> bool {
        if matches!(change.event_type, chainfire_client::EventType::Delete) {
            self.watched_instance_signatures.remove(key);
            return true;
        }

        let Some(signature) = canonical_json_signature(&change.value, &["last_heartbeat", "observed_at"])
        else {
            return true;
        };
        if self.watched_instance_signatures.get(key) == Some(&signature) {
            return false;
        }
        self.watched_instance_signatures
            .insert(key.to_string(), signature);
        true
    }

    fn apply_deployment_spec_watch_change(&mut self, key: &str, change: &WatchChange) -> bool {
        if matches!(change.event_type, chainfire_client::EventType::Delete) {
            self.watched_deployment_spec_signatures.remove(key);
            return true;
        }

        let Some(signature) = canonical_json_signature(&change.value, &[]) else {
            return true;
        };
        if self.watched_deployment_spec_signatures.get(key) == Some(&signature) {
            return false;
        }
        self.watched_deployment_spec_signatures
            .insert(key.to_string(), signature);
        true
    }

    fn apply_deployment_status_watch_change(&mut self, key: &str, change: &WatchChange) -> bool {
        if self.consume_pending_delete(key, change) {
            self.watched_deployment_status_signatures.remove(key);
            return false;
        }
        if let Some(payload) = self.consume_pending_write(key, change) {
            if let Some(signature) = host_deployment_status_signature(&payload) {
                self.watched_deployment_status_signatures
                    .insert(key.to_string(), signature);
            } else {
                self.watched_deployment_status_signatures.remove(key);
            }
            return false;
        }
        if matches!(change.event_type, chainfire_client::EventType::Delete) {
            self.watched_deployment_status_signatures.remove(key);
            return true;
        }

        let Some(signature) = host_deployment_status_signature(&change.value) else {
            return true;
        };
        if self.watched_deployment_status_signatures.get(key) == Some(&signature) {
            return false;
        }
        self.watched_deployment_status_signatures
            .insert(key.to_string(), signature);
        true
    }

    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 aborted = existing_status
        .and_then(|status| status.phase.as_deref())
        == Some("aborted");
    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<_>>();

    if aborted {
        desired_deletes = desired_systems
            .iter()
            .filter(|(_, desired)| {
                desired.deployment_id.as_deref() == Some(deployment.name.as_str())
            })
            .map(|(node_id, _)| node_id.clone())
            .collect();

        for node in &selector_matches {
            if 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);
            }
        }

        let mut status = existing_status.cloned().unwrap_or_default();
        status.name = deployment.name.clone();
        status.paused = true;
        status.paused_by_operator = true;
        status.phase = Some("aborted".to_string());
        if status.message.is_none() {
            status.message = Some("aborted by operator".to_string());
        }
        status.updated_at = Some(now);

        desired_deletes.sort();
        desired_deletes.dedup();

        return HostDeploymentPlan {
            status,
            desired_upserts: Vec::new(),
            desired_deletes,
            node_updates,
        };
    }

    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::*;
    use chainfire_client::EventType;

    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/ultracloud-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),
        }
    }

    fn test_controller() -> HostDeploymentController {
        HostDeploymentController::new(HostsCommand {
            endpoint: "http://127.0.0.1:7000".to_string(),
            cluster_namespace: "ultracloud".to_string(),
            cluster_id: "test-cluster".to_string(),
            interval_secs: 1,
            heartbeat_timeout_secs: 300,
            dry_run: true,
            once: 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()]);
    }

    #[test]
    fn plan_rollout_keeps_aborted_state_and_clears_desired_system() {
        let deployment = test_deployment();
        let mut draining = test_node("node01", "rack-a");
        draining.state = Some("draining".to_string());
        let nodes = vec![draining];
        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 status = HostDeploymentStatus {
            name: "worker-rollout".to_string(),
            phase: Some("aborted".to_string()),
            paused: true,
            paused_by_operator: true,
            message: Some("aborted by operator".to_string()),
            ..HostDeploymentStatus::default()
        };

        let plan = plan_host_deployment(
            &deployment,
            Some(&status),
            &nodes,
            &desired,
            &HashMap::new(),
            &[],
            300,
        );

        assert!(plan.desired_upserts.is_empty());
        assert_eq!(plan.desired_deletes, vec!["node01".to_string()]);
        assert_eq!(plan.status.phase.as_deref(), Some("aborted"));
        assert!(plan.status.paused);
        assert_eq!(
            plan.node_updates
                .get("node01")
                .and_then(|node| node.state.as_deref()),
            Some("active")
        );
    }

    #[test]
    fn host_watch_ignores_node_heartbeat_only_changes() {
        let mut controller = test_controller();
        let node_key = key_node(
            &controller.cluster_namespace,
            &controller.cluster_id,
            "node01",
        );

        let first = WatchChange {
            event_type: EventType::Put,
            key: node_key.clone(),
            value: serde_json::to_vec(&serde_json::json!({
                "node_id": "node01",
                "hostname": "node01",
                "ip": "10.0.0.11",
                "roles": ["worker"],
                "labels": {"tier": "general"},
                "state": "active",
                "last_heartbeat": "2026-04-01T00:00:00Z"
            }))
            .unwrap(),
            revision: 1,
        };
        assert!(controller.apply_watch_change(&first));

        let heartbeat_only = WatchChange {
            revision: 2,
            value: serde_json::to_vec(&serde_json::json!({
                "node_id": "node01",
                "hostname": "node01",
                "ip": "10.0.0.11",
                "roles": ["worker"],
                "labels": {"tier": "general"},
                "state": "active",
                "last_heartbeat": "2026-04-01T00:00:05Z"
            }))
            .unwrap(),
            ..first.clone()
        };
        assert!(!controller.apply_watch_change(&heartbeat_only));

        let drain = WatchChange {
            revision: 3,
            value: serde_json::to_vec(&serde_json::json!({
                "node_id": "node01",
                "hostname": "node01",
                "ip": "10.0.0.11",
                "roles": ["worker"],
                "labels": {"tier": "general"},
                "state": "draining",
                "last_heartbeat": "2026-04-01T00:00:10Z"
            }))
            .unwrap(),
            ..first
        };
        assert!(controller.apply_watch_change(&drain));
    }

    #[test]
    fn host_watch_reacts_only_to_operator_pause_or_abort_status_changes() {
        let mut controller = test_controller();
        let key = key_host_deployment_status(
            &controller.cluster_namespace,
            &controller.cluster_id,
            "worker-rollout",
        );

        let running = WatchChange {
            event_type: EventType::Put,
            key: key.clone(),
            value: serde_json::to_vec(&serde_json::json!({
                "name": "worker-rollout",
                "phase": "running",
                "paused": false,
                "paused_by_operator": false,
                "selected_nodes": ["node01"],
                "in_progress_nodes": ["node01"],
                "updated_at": "2026-04-01T00:00:00Z"
            }))
            .unwrap(),
            revision: 1,
        };
        assert!(controller.apply_watch_change(&running));

        let completed = WatchChange {
            revision: 2,
            value: serde_json::to_vec(&serde_json::json!({
                "name": "worker-rollout",
                "phase": "completed",
                "paused": false,
                "paused_by_operator": false,
                "selected_nodes": ["node01"],
                "completed_nodes": ["node01"],
                "updated_at": "2026-04-01T00:00:10Z"
            }))
            .unwrap(),
            ..running.clone()
        };
        assert!(!controller.apply_watch_change(&completed));

        let paused = WatchChange {
            revision: 3,
            value: serde_json::to_vec(&serde_json::json!({
                "name": "worker-rollout",
                "phase": "paused",
                "paused": true,
                "paused_by_operator": true,
                "updated_at": "2026-04-01T00:00:20Z"
            }))
            .unwrap(),
            ..running.clone()
        };
        assert!(controller.apply_watch_change(&paused));

        let aborted = WatchChange {
            revision: 4,
            value: serde_json::to_vec(&serde_json::json!({
                "name": "worker-rollout",
                "phase": "aborted",
                "paused": true,
                "paused_by_operator": true,
                "updated_at": "2026-04-01T00:00:30Z"
            }))
            .unwrap(),
            ..running
        };
        assert!(controller.apply_watch_change(&aborted));
    }
}
