use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::PathBuf;
use std::process::Stdio;
use std::time::Duration;

use anyhow::{Context, Result};
use chainfire_client::{Client, ClientError};
use chrono::{DateTime, Utc};
use deployer_types::{ContainerSpec, HealthCheckSpec, ProcessSpec, ServiceInstanceSpec};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::process::Command;
use tokio::time::sleep;
use tracing::{info, warn};

use crate::process::{render_container_process_spec, ProcessManager};

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_instance(
    cluster_namespace: &str,
    cluster_id: &str,
    service: &str,
    instance_id: &str,
) -> Vec<u8> {
    format!(
        "{}instances/{}/{}",
        cluster_prefix(cluster_namespace, cluster_id),
        service,
        instance_id
    )
    .into_bytes()
}

#[derive(Debug, Deserialize, Serialize)]
pub struct NodeState {
    pub node_id: String,
    pub ip: String,
    pub hostname: String,
    #[serde(default)]
    pub roles: Vec<String>,
    #[serde(default)]
    pub labels: std::collections::HashMap<String, String>,
    #[serde(default)]
    pub state: Option<String>,
    #[serde(default)]
    pub last_heartbeat: Option<DateTime<Utc>>,
    #[serde(flatten)]
    pub extra: std::collections::HashMap<String, Value>,
}

pub struct Agent {
    endpoint: String,
    cluster_namespace: String,
    cluster_id: String,
    node_id: String,
    interval: Duration,
    apply: bool,
    allow_local_instance_upsert: bool,
    process_manager: ProcessManager,
    next_health_checks: HashMap<String, DateTime<Utc>>,
}

#[derive(Debug, Deserialize, Serialize, Clone)]
struct LocalInstanceSpec {
    service: String,
    instance_id: String,
    #[serde(default)]
    node_id: Option<String>,
    ip: String,
    port: u16,
    #[serde(default)]
    mesh_port: Option<u16>,
    #[serde(default)]
    health_check: Option<HealthCheckSpec>,
    #[serde(default)]
    process: Option<ProcessSpec>,
    #[serde(default)]
    container: Option<ContainerSpec>,
}

impl Agent {
    pub fn new(
        endpoint: String,
        cluster_namespace: String,
        cluster_id: String,
        node_id: String,
        interval: Duration,
        apply: bool,
        allow_local_instance_upsert: bool,
        pid_dir: PathBuf,
    ) -> Self {
        Self {
            endpoint,
            cluster_namespace,
            cluster_id,
            node_id,
            interval,
            apply,
            allow_local_instance_upsert,
            process_manager: ProcessManager::new(pid_dir),
            next_health_checks: HashMap::new(),
        }
    }

    pub async fn run_loop(&mut self) -> Result<()> {
        loop {
            if let Err(e) = self.tick().await {
                warn!(error = %e, "node-agent tick failed");
            }
            sleep(self.interval).await;
        }
    }

    pub async fn tick(&mut self) -> Result<()> {
        let mut client = Client::connect(self.endpoint.clone()).await?;

        // Node 情報
        let node_key = key_node(&self.cluster_namespace, &self.cluster_id, &self.node_id);
        let node_raw = client.get_with_revision(&node_key).await?;
        let Some((node_bytes, node_revision)) = node_raw else {
            warn!(
                "node definition not found in Chainfire for cluster_id={}, node_id={}",
                self.cluster_id, self.node_id
            );
            return Ok(());
        };

        let mut node: NodeState = match serde_json::from_slice(&node_bytes) {
            Ok(n) => n,
            Err(e) => {
                warn!(error = %e, "failed to parse node JSON");
                return Ok(());
            }
        };

        // Heartbeat を更新し、Chainfire 上の Node を upsert
        node.last_heartbeat = Some(Utc::now());
        if self.apply {
            let should_activate = match node.state.as_deref() {
                Some("failed") | Some("draining") => false,
                Some("active") => false,
                Some("provisioning") | Some("pending") | None => true,
                Some(other) => {
                    warn!(state = %other, "unknown node state; leaving unchanged");
                    false
                }
            };
            if should_activate {
                node.state = Some("active".to_string());
            }
        }
        let updated = serde_json::to_vec(&node)?;
        if let Err(e) = client
            .put_if_revision(&node_key, &updated, node_revision)
            .await
        {
            warn!(error = %e, "failed to update node heartbeat");
        }

        // ローカル定義された ServiceInstance を Chainfire に登録
        if self.allow_local_instance_upsert {
            if let Err(e) = self.sync_local_instances(&mut client).await {
                warn!(error = %e, "failed to sync local service instances");
            }
        } else {
            info!("local instance upsert disabled; skipping /etc/photoncloud/instances.json");
        }

        if self.apply {
            // プロセスの起動/停止をReconcile
            if let Err(e) = self.reconcile_processes(&mut client).await {
                warn!(error = %e, "failed to reconcile processes");
            }

            // ヘルスチェックを実行して状態を更新
            if let Err(e) = self.update_health_status(&mut client).await {
                warn!(error = %e, "failed to update health status");
            }
        } else {
            info!("apply=false; skipping process reconciliation and health checks");
        }

        self.log_node_only(&node);

        Ok(())
    }

    fn log_node_only(&self, node: &NodeState) {
        info!(
            node_id = %node.node_id,
            hostname = %node.hostname,
            roles = ?node.roles,
            "observed desired state for node"
        );
    }

    fn is_local_instance(&self, inst: &LocalInstanceSpec) -> bool {
        matches!(inst.node_id.as_deref(), Some(node_id) if node_id == self.node_id)
    }

    fn is_local_managed_instance(&self, inst: &ServiceInstanceSpec) -> bool {
        inst.node_id == self.node_id
    }

    fn render_template_value(&self, value: &str, inst: &ServiceInstanceSpec) -> String {
        let mesh_port = inst
            .mesh_port
            .map(|port| port.to_string())
            .unwrap_or_default();
        [
            ("${SERVICE}", inst.service.as_str()),
            ("${INSTANCE_ID}", inst.instance_id.as_str()),
            ("${NODE_ID}", inst.node_id.as_str()),
            ("${INSTANCE_IP}", inst.ip.as_str()),
            ("${INSTANCE_PORT}", &inst.port.to_string()),
            ("${MESH_PORT}", mesh_port.as_str()),
        ]
        .into_iter()
        .fold(value.to_string(), |acc, (needle, replacement)| {
            acc.replace(needle, replacement)
        })
    }

    fn render_process_spec(&self, spec: &ProcessSpec, inst: &ServiceInstanceSpec) -> ProcessSpec {
        ProcessSpec {
            command: self.render_template_value(&spec.command, inst),
            args: spec
                .args
                .iter()
                .map(|arg| self.render_template_value(arg, inst))
                .collect(),
            working_dir: spec
                .working_dir
                .as_deref()
                .map(|value| self.render_template_value(value, inst)),
            env: spec
                .env
                .iter()
                .map(|(key, value)| (key.clone(), self.render_template_value(value, inst)))
                .collect(),
        }
    }

    fn render_container_spec(
        &self,
        spec: &ContainerSpec,
        inst: &ServiceInstanceSpec,
    ) -> ContainerSpec {
        let mut rendered = spec.clone();
        rendered.image = self.render_template_value(&rendered.image, inst);
        rendered.command = rendered
            .command
            .iter()
            .map(|arg| self.render_template_value(arg, inst))
            .collect();
        rendered.args = rendered
            .args
            .iter()
            .map(|arg| self.render_template_value(arg, inst))
            .collect();
        rendered.working_dir = rendered
            .working_dir
            .as_deref()
            .map(|value| self.render_template_value(value, inst));
        rendered.env = rendered
            .env
            .iter()
            .map(|(key, value)| (key.clone(), self.render_template_value(value, inst)))
            .collect();
        for volume in &mut rendered.volumes {
            volume.source = self.render_template_value(&volume.source, inst);
            volume.target = self.render_template_value(&volume.target, inst);
        }
        rendered
    }

    fn desired_process_spec(&self, inst: &ServiceInstanceSpec) -> Option<ProcessSpec> {
        match (&inst.container, &inst.process) {
            (Some(container), maybe_process) => {
                if maybe_process.is_some() {
                    warn!(
                        service = %inst.service,
                        instance_id = %inst.instance_id,
                        "instance defines both process and container; container takes precedence"
                    );
                }
                let rendered = self.render_container_spec(container, inst);
                Some(render_container_process_spec(
                    &inst.service,
                    &inst.instance_id,
                    inst.port,
                    &rendered,
                ))
            }
            (None, Some(process)) => Some(self.render_process_spec(process, inst)),
            (None, None) => None,
        }
    }

    fn apply_instance_health_fields(
        inst_value: &mut Value,
        started_at: &DateTime<Utc>,
        heartbeat_at: &DateTime<Utc>,
        health_status: &str,
    ) {
        let Some(obj) = inst_value.as_object_mut() else {
            return;
        };

        let observed_at = Value::String(started_at.to_rfc3339());
        match obj.get_mut("observed_at") {
            Some(slot) if slot.is_null() => *slot = observed_at,
            Some(_) => {}
            None => {
                obj.insert("observed_at".to_string(), observed_at);
            }
        }

        obj.insert(
            "state".to_string(),
            Value::String(health_status.to_string()),
        );
        obj.insert(
            "last_heartbeat".to_string(),
            Value::String(heartbeat_at.to_rfc3339()),
        );
    }

    async fn persist_instance_health(
        &self,
        client: &mut Client,
        key: &[u8],
        inst: &ServiceInstanceSpec,
        mut inst_value: Value,
        mut mod_revision: u64,
        started_at: &DateTime<Utc>,
        health_status: &str,
    ) -> Result<()> {
        for attempt in 0..3 {
            let heartbeat_at = Utc::now();
            Self::apply_instance_health_fields(
                &mut inst_value,
                started_at,
                &heartbeat_at,
                health_status,
            );

            let updated = serde_json::to_vec(&inst_value)?;
            match client.put_if_revision(key, &updated, mod_revision).await {
                Ok(()) => return Ok(()),
                Err(ClientError::Conflict(error)) if attempt < 2 => {
                    warn!(
                        service = %inst.service,
                        instance_id = %inst.instance_id,
                        mod_revision,
                        attempt = attempt + 1,
                        error = %error,
                        "instance health update raced with another writer; retrying with fresh revision"
                    );

                    let Some((latest_bytes, latest_revision)) =
                        client.get_with_revision(key).await?
                    else {
                        warn!(
                            service = %inst.service,
                            instance_id = %inst.instance_id,
                            "instance disappeared while retrying health update"
                        );
                        return Ok(());
                    };

                    inst_value = serde_json::from_slice(&latest_bytes).with_context(|| {
                        format!(
                            "failed to parse refreshed instance JSON for {}",
                            inst.instance_id
                        )
                    })?;
                    mod_revision = latest_revision;
                }
                Err(error) => return Err(error.into()),
            }
        }

        Ok(())
    }

    /// ローカルファイル (/etc/photoncloud/instances.json) から ServiceInstance 定義を読み、
    /// Chainfire 上の `photoncloud/clusters/{cluster_id}/instances/{service}/{instance_id}` に upsert する。
    async fn sync_local_instances(&self, client: &mut Client) -> Result<()> {
        let path = PathBuf::from("/etc/photoncloud/instances.json");
        let contents = match fs::read_to_string(&path) {
            Ok(c) => c,
            Err(e) => {
                warn!(error = %e, "no local instances file found, skipping");
                return Ok(());
            }
        };

        let instances: Vec<LocalInstanceSpec> = serde_json::from_str(&contents)
            .with_context(|| format!("failed to parse {}", path.display()))?;

        for inst in &instances {
            let mut inst = inst.clone();
            if inst.node_id.is_none() {
                inst.node_id = Some(self.node_id.clone());
            }
            if !self.is_local_instance(&inst) {
                continue;
            }

            let key = key_instance(
                &self.cluster_namespace,
                &self.cluster_id,
                &inst.service,
                &inst.instance_id,
            );
            let mut desired_value = serde_json::to_value(&inst)?;

            if let Some(existing) = client.get(&key).await? {
                if let Ok(existing_value) = serde_json::from_slice::<Value>(&existing) {
                    if let (Some(existing_obj), Some(desired_obj)) =
                        (existing_value.as_object(), desired_value.as_object_mut())
                    {
                        for preserve_key in ["state", "last_heartbeat", "observed_at"] {
                            if let Some(value) = existing_obj.get(preserve_key) {
                                desired_obj
                                    .entry(preserve_key.to_string())
                                    .or_insert(value.clone());
                            }
                        }
                    }
                }
            }

            let value = serde_json::to_vec(&desired_value)?;
            client.put(&key, &value).await?;
            info!(
                service = %inst.service,
                instance_id = %inst.instance_id,
                "synced local ServiceInstance to Chainfire"
            );
        }

        Ok(())
    }

    /// Desired Stateに基づいてプロセスを起動/停止する
    async fn reconcile_processes(&mut self, client: &mut Client) -> Result<()> {
        let prefix = format!(
            "{}instances/",
            cluster_prefix(&self.cluster_namespace, &self.cluster_id)
        );
        let (kvs, _) = client.scan_prefix(prefix.as_bytes(), 0).await?;

        let mut desired_instances = Vec::new();
        let mut desired_keys = HashSet::new();
        for (_key, value, _) in kvs {
            let inst: ServiceInstanceSpec = match serde_json::from_slice(&value) {
                Ok(i) => i,
                Err(e) => {
                    warn!(error = %e, "failed to parse instance");
                    continue;
                }
            };

            if !self.is_local_managed_instance(&inst) {
                continue;
            }

            if let Some(proc_spec) = self.desired_process_spec(&inst) {
                desired_keys.insert(format!("{}/{}", inst.service, inst.instance_id));
                desired_instances.push((inst.service.clone(), inst.instance_id.clone(), proc_spec));
            }
        }

        // Desired Stateに基づいてプロセスを管理
        for (service, instance_id, proc_spec) in desired_instances {
            if let Some(process) = self.process_manager.get_mut(&service, &instance_id) {
                if process.spec != proc_spec {
                    process.spec = proc_spec;
                    info!(
                        service = %service,
                        instance_id = %instance_id,
                        "process spec changed, restarting"
                    );
                    if let Err(e) = process.restart().await {
                        warn!(
                            service = %service,
                            instance_id = %instance_id,
                            error = ?e,
                            "failed to restart process after spec update"
                        );
                    }
                }
            } else {
                // 新しいプロセスを追加
                self.process_manager
                    .add(service.clone(), instance_id.clone(), proc_spec);
                info!(
                    service = %service,
                    instance_id = %instance_id,
                    "added new process to manager"
                );
            }
        }

        self.process_manager.stop_unmanaged(&desired_keys).await?;

        // Desired Stateにないプロセスは停止して削除
        for (service, instance_id) in self.process_manager.list_instances() {
            let key = format!("{}/{}", service, instance_id);
            if !desired_keys.contains(&key) {
                if let Err(e) = self
                    .process_manager
                    .stop_and_remove(&service, &instance_id)
                    .await
                {
                    warn!(
                        service = %service,
                        instance_id = %instance_id,
                        error = %e,
                        "failed to stop/remove process"
                    );
                }
            }
        }

        // Reconcile: 停止しているプロセスを再起動
        self.process_manager.reconcile().await?;

        Ok(())
    }

    /// 各ServiceInstanceのヘルスチェックを実行し、Chainfire上の状態を更新
    fn health_check_key(inst: &ServiceInstanceSpec) -> String {
        format!("{}/{}", inst.service, inst.instance_id)
    }

    fn health_check_due(&self, inst: &ServiceInstanceSpec, _spec: &HealthCheckSpec) -> bool {
        let key = Self::health_check_key(inst);
        match self.next_health_checks.get(&key) {
            Some(next_due) => Utc::now() >= *next_due,
            None => true,
        }
    }

    fn schedule_next_health_check(&mut self, inst: &ServiceInstanceSpec, spec: &HealthCheckSpec) {
        let interval_secs = spec.interval_secs.unwrap_or(self.interval.as_secs()).max(1);
        let key = Self::health_check_key(inst);
        self.next_health_checks.insert(
            key,
            Utc::now() + chrono::Duration::seconds(interval_secs as i64),
        );
    }

    fn startup_grace_secs(&self, inst: &ServiceInstanceSpec, spec: &HealthCheckSpec) -> u64 {
        spec.startup_grace_secs.unwrap_or_else(|| {
            if inst.container.is_some() {
                let interval = spec.interval_secs.unwrap_or(self.interval.as_secs()).max(1);
                let timeout = spec.timeout_secs.unwrap_or(5).max(1);
                std::cmp::max(60, std::cmp::max(interval * 6, timeout * 6))
            } else {
                0
            }
        })
    }

    async fn process_running(&mut self, inst: &ServiceInstanceSpec) -> bool {
        self.process_manager
            .is_running(&inst.service, &inst.instance_id)
            .await
            .unwrap_or(false)
    }

    fn starting_within_grace(
        &self,
        inst: &ServiceInstanceSpec,
        spec: &HealthCheckSpec,
        process_running: bool,
        started_at: DateTime<Utc>,
        now: DateTime<Utc>,
        health_status: &str,
    ) -> bool {
        if health_status == "healthy" || !process_running {
            return false;
        }

        let startup_grace_secs = self.startup_grace_secs(inst, spec);
        if startup_grace_secs == 0 {
            return false;
        }

        now.signed_duration_since(started_at).num_seconds() < startup_grace_secs as i64
    }

    async fn update_health_status(&mut self, client: &mut Client) -> Result<()> {
        let prefix = format!(
            "{}instances/",
            cluster_prefix(&self.cluster_namespace, &self.cluster_id)
        );
        let (kvs, _) = client.scan_prefix(prefix.as_bytes(), 0).await?;
        let mut seen = HashSet::new();

        for (key, value, mod_revision) in kvs {
            let inst_value: Value = match serde_json::from_slice(&value) {
                Ok(v) => v,
                Err(e) => {
                    warn!(error = %e, "failed to parse instance json");
                    continue;
                }
            };

            let inst: ServiceInstanceSpec = match serde_json::from_value(inst_value.clone()) {
                Ok(i) => i,
                Err(e) => {
                    warn!(error = %e, "failed to parse instance");
                    continue;
                }
            };

            if !self.is_local_managed_instance(&inst) {
                continue;
            }

            let health_key = Self::health_check_key(&inst);
            seen.insert(health_key.clone());

            if let Some(ref health_check) = inst.health_check {
                if !self.health_check_due(&inst, health_check) {
                    continue;
                }
            }

            let now = Utc::now();
            let process_running = self.process_running(&inst).await;
            let started_at = self
                .process_manager
                .started_at(&inst.service, &inst.instance_id)
                .or(inst.observed_at)
                .unwrap_or(now);

            let health_status = if let Some(ref health_check) = inst.health_check {
                let status = self.check_health(&inst, health_check).await;
                self.schedule_next_health_check(&inst, health_check);
                if self.starting_within_grace(
                    &inst,
                    health_check,
                    process_running,
                    started_at,
                    now,
                    &status,
                ) {
                    "starting".to_string()
                } else {
                    status
                }
            } else {
                "healthy".to_string() // デフォルトはhealthy
            };

            if let Err(e) = self
                .persist_instance_health(
                    client,
                    &key,
                    &inst,
                    inst_value,
                    mod_revision,
                    &started_at,
                    &health_status,
                )
                .await
            {
                warn!(
                    service = %inst.service,
                    instance_id = %inst.instance_id,
                    mod_revision,
                    error = ?e,
                    "failed to update instance health status"
                );
            }

            info!(
                service = %inst.service,
                instance_id = %inst.instance_id,
                status = %health_status,
                "health check completed"
            );
        }

        self.next_health_checks.retain(|key, _| seen.contains(key));

        Ok(())
    }

    async fn check_health(&self, inst: &ServiceInstanceSpec, spec: &HealthCheckSpec) -> String {
        match spec.check_type.as_str() {
            "http" => {
                if let Some(ref path) = spec.path {
                    let url = format!("http://{}:{}{}", inst.ip, inst.port, path);
                    let client = reqwest::Client::builder()
                        .timeout(Duration::from_secs(spec.timeout_secs.unwrap_or(5)))
                        .build();
                    if let Ok(c) = client {
                        if let Ok(resp) = c.get(&url).send().await {
                            if resp.status().is_success() {
                                return "healthy".to_string();
                            }
                        }
                    }
                }
                "unhealthy".to_string()
            }
            "tcp" => {
                use tokio::net::TcpStream;
                let addr = format!("{}:{}", inst.ip, inst.port);
                match tokio::time::timeout(
                    Duration::from_secs(spec.timeout_secs.unwrap_or(5)),
                    TcpStream::connect(&addr),
                )
                .await
                {
                    Ok(Ok(_)) => "healthy".to_string(),
                    _ => "unhealthy".to_string(),
                }
            }
            "command" => {
                if let Some(ref cmd) = spec.path {
                    match Command::new("sh")
                        .arg("-c")
                        .arg(cmd)
                        .stdout(Stdio::null())
                        .stderr(Stdio::null())
                        .status()
                        .await
                    {
                        Ok(status) if status.success() => "healthy".to_string(),
                        _ => "unhealthy".to_string(),
                    }
                } else {
                    "unknown".to_string()
                }
            }
            _ => "unknown".to_string(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use std::path::PathBuf;

    fn test_agent() -> Agent {
        Agent::new(
            "http://127.0.0.1:7000".to_string(),
            "photoncloud".to_string(),
            "test-cluster".to_string(),
            "node01".to_string(),
            Duration::from_secs(1),
            false,
            false,
            PathBuf::from("/tmp/photoncloud-node-agent-tests"),
        )
    }

    fn test_instance() -> ServiceInstanceSpec {
        ServiceInstanceSpec {
            instance_id: "api-node01".to_string(),
            service: "api".to_string(),
            node_id: "node01".to_string(),
            ip: "127.0.0.2".to_string(),
            port: 18080,
            mesh_port: Some(28080),
            version: None,
            health_check: None,
            process: Some(ProcessSpec {
                command: "python3".to_string(),
                args: vec![
                    "-m".to_string(),
                    "http.server".to_string(),
                    "${INSTANCE_PORT}".to_string(),
                    "--bind".to_string(),
                    "${INSTANCE_IP}".to_string(),
                ],
                working_dir: Some("/srv/${SERVICE}".to_string()),
                env: HashMap::from([("INSTANCE".to_string(), "${INSTANCE_ID}".to_string())]),
            }),
            container: None,
            managed_by: None,
            state: None,
            last_heartbeat: None,
            observed_at: None,
        }
    }

    #[test]
    fn test_render_process_spec_templates() {
        let agent = test_agent();
        let instance = test_instance();
        let rendered = agent.render_process_spec(instance.process.as_ref().unwrap(), &instance);

        assert_eq!(rendered.args[2], "18080");
        assert_eq!(rendered.args[4], "127.0.0.2");
        assert_eq!(rendered.working_dir.as_deref(), Some("/srv/api"));
        assert_eq!(
            rendered.env.get("INSTANCE").map(String::as_str),
            Some("api-node01")
        );
    }

    #[test]
    fn test_container_takes_precedence_when_present() {
        let agent = test_agent();
        let mut instance = test_instance();
        instance.container = Some(ContainerSpec {
            image: "nginx:latest".to_string(),
            runtime: Some("docker".to_string()),
            command: vec![],
            args: vec![],
            env: HashMap::new(),
            ports: vec![],
            volumes: vec![],
            network_mode: None,
            pull_policy: None,
            working_dir: None,
        });

        let rendered = agent.desired_process_spec(&instance).unwrap();
        assert_eq!(rendered.command, "docker");
    }

    #[test]
    fn test_health_check_interval_is_tracked_per_instance() {
        let mut agent = test_agent();
        let mut instance = test_instance();
        let health_check = HealthCheckSpec {
            check_type: "http".to_string(),
            path: Some("/health".to_string()),
            interval_secs: Some(30),
            timeout_secs: Some(5),
            startup_grace_secs: Some(30),
        };
        instance.health_check = Some(health_check.clone());

        assert!(agent.health_check_due(&instance, &health_check));
        agent.schedule_next_health_check(&instance, &health_check);
        assert!(!agent.health_check_due(&instance, &health_check));

        let key = Agent::health_check_key(&instance);
        agent
            .next_health_checks
            .insert(key, Utc::now() - chrono::Duration::seconds(1));
        assert!(agent.health_check_due(&instance, &health_check));
    }

    #[test]
    fn test_apply_instance_health_fields_replaces_nulls_and_preserves_observed_at() {
        let started_at = DateTime::parse_from_rfc3339("2026-03-31T03:00:00Z")
            .unwrap()
            .with_timezone(&Utc);
        let heartbeat_at = DateTime::parse_from_rfc3339("2026-03-31T03:00:05Z")
            .unwrap()
            .with_timezone(&Utc);
        let original_observed_at = "2026-03-31T02:59:59Z";

        let mut null_observed = serde_json::json!({
            "observed_at": null,
            "state": null,
            "last_heartbeat": null
        });
        Agent::apply_instance_health_fields(
            &mut null_observed,
            &started_at,
            &heartbeat_at,
            "healthy",
        );
        assert_eq!(
            null_observed.get("observed_at").and_then(Value::as_str),
            Some("2026-03-31T03:00:00+00:00")
        );
        assert_eq!(
            null_observed.get("state").and_then(Value::as_str),
            Some("healthy")
        );
        assert_eq!(
            null_observed.get("last_heartbeat").and_then(Value::as_str),
            Some("2026-03-31T03:00:05+00:00")
        );

        let mut existing_observed = serde_json::json!({
            "observed_at": original_observed_at,
            "state": "starting",
            "last_heartbeat": null
        });
        Agent::apply_instance_health_fields(
            &mut existing_observed,
            &started_at,
            &heartbeat_at,
            "healthy",
        );
        assert_eq!(
            existing_observed.get("observed_at").and_then(Value::as_str),
            Some(original_observed_at)
        );
        assert_eq!(
            existing_observed.get("state").and_then(Value::as_str),
            Some("healthy")
        );
    }
}
