use std::sync::Arc;

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::http::{HeaderMap, Uri};
use axum::response::Response;
use chrono::{DateTime, Duration, Utc};
use futures_util::{SinkExt, StreamExt};
use serde_json::{json, Value};
use tokio::sync::mpsc;

use crate::db::{
    raw_task, record_event, update_provider_heartbeat, upsert_docker_status, upsert_provider_node,
    upsert_system_status,
};
use crate::egress_tcp::{handle_egress_tcp_close, handle_egress_tcp_data, handle_egress_tcp_open};
use crate::http::json_response;
use crate::json_util::compact_json_for_storage;
use crate::ssh_bridge::forward_ssh_provider_message;
use crate::state::{AppState, HttpTunnelResponse, ProviderConnection};
use crate::task_dispatcher::{is_terminal_task_status, notify_task_terminal};
use crate::types::RawTaskRow;

const PROVIDER_UPGRADE_FINAL_HEARTBEAT_TIMEOUT_MS: i64 = 240_000;
const PROVIDER_UPGRADE_FINAL_HEARTBEAT_POLL_MS: u64 = 2_000;

#[derive(Clone, Debug)]
struct ProviderUpgradeFinalHeartbeatExpectation {
    task_id: String,
    provider_id: String,
    target_version: String,
    min_heartbeat_at: DateTime<Utc>,
    deadline_at: DateTime<Utc>,
    scheduler_result: Value,
    candidate_name: Option<String>,
    last_known_good_version: Option<String>,
}

#[derive(Clone, Debug)]
struct ProviderUpgradeNodeSnapshot {
    provider_id: String,
    status: String,
    labels: Value,
}

#[derive(Clone, Debug)]
enum ProviderUpgradeFinalHeartbeatDecision {
    Pending(Value),
    Succeeded(Value),
    Failed(Value),
}

fn json_string_at(value: &Value, path: &[&str]) -> Option<String> {
    let mut current = value;
    for part in path {
        current = current.get(*part)?;
    }
    current
        .as_str()
        .filter(|text| !text.is_empty())
        .map(ToOwned::to_owned)
}

fn provider_upgrade_label_string(labels: &Value, key: &str) -> Option<String> {
    labels
        .get(key)
        .and_then(Value::as_str)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned)
}

fn provider_upgrade_heartbeat_at(snapshot: &ProviderUpgradeNodeSnapshot) -> Option<DateTime<Utc>> {
    provider_upgrade_label_string(&snapshot.labels, "providerGatewayHeartbeatAt").and_then(
        |value| {
            DateTime::parse_from_rfc3339(&value)
                .ok()
                .map(|parsed| parsed.with_timezone(&Utc))
        },
    )
}

fn provider_upgrade_schedule_task(task: &RawTaskRow) -> bool {
    task.command == "provider.upgrade"
        && task.payload.get("mode").and_then(Value::as_str) == Some("schedule")
}

fn provider_upgrade_target_version(task: &RawTaskRow, scheduler_result: &Value) -> Option<String> {
    json_string_at(scheduler_result, &["targetProviderGatewayVersion"])
        .or_else(|| json_string_at(scheduler_result, &["plan", "targetProviderGatewayVersion"]))
        .or_else(|| json_string_at(&task.payload, &["targetProviderGatewayVersion"]))
}

fn provider_upgrade_last_known_good_version(scheduler_result: &Value) -> Option<String> {
    json_string_at(scheduler_result, &["providerGatewayVersion"])
        .or_else(|| json_string_at(scheduler_result, &["plan", "providerGatewayVersion"]))
}

fn provider_upgrade_candidate_name(scheduler_result: &Value) -> Option<String> {
    json_string_at(scheduler_result, &["candidateName"])
        .or_else(|| json_string_at(scheduler_result, &["plan", "candidateName"]))
}

fn provider_upgrade_expectation(
    task: &RawTaskRow,
    scheduler_result: &Value,
    now: DateTime<Utc>,
) -> Option<ProviderUpgradeFinalHeartbeatExpectation> {
    let target_version = provider_upgrade_target_version(task, scheduler_result)?;
    Some(ProviderUpgradeFinalHeartbeatExpectation {
        task_id: task.id.clone(),
        provider_id: task.provider_id.clone(),
        target_version,
        min_heartbeat_at: now,
        deadline_at: now + Duration::milliseconds(PROVIDER_UPGRADE_FINAL_HEARTBEAT_TIMEOUT_MS),
        scheduler_result: scheduler_result.clone(),
        candidate_name: provider_upgrade_candidate_name(scheduler_result),
        last_known_good_version: provider_upgrade_last_known_good_version(scheduler_result),
    })
}

fn provider_upgrade_snapshot_is_candidate(
    snapshot: &ProviderUpgradeNodeSnapshot,
    candidate_name: Option<&str>,
) -> bool {
    let container_name =
        provider_upgrade_label_string(&snapshot.labels, "providerGatewayContainerName");
    if let (Some(container_name), Some(candidate_name)) =
        (container_name.as_deref(), candidate_name)
    {
        if container_name == candidate_name {
            return true;
        }
    }
    container_name
        .as_deref()
        .is_some_and(|name| name.contains("-candidate-"))
}

fn insert_json_field_with_unavailable_reason(
    object: &mut serde_json::Map<String, Value>,
    key: &str,
    value: Option<String>,
    unavailable_reason: &str,
) {
    if let Some(value) = value {
        object.insert(key.to_string(), Value::String(value));
    } else {
        object.insert(key.to_string(), Value::Null);
        object.insert(
            format!("{key}UnavailableReason"),
            Value::String(unavailable_reason.to_string()),
        );
    }
}

fn provider_upgrade_final_container_json(
    snapshot: Option<&ProviderUpgradeNodeSnapshot>,
    unavailable_reason: &str,
) -> Value {
    let labels = snapshot.map(|snapshot| &snapshot.labels);
    let heartbeat = snapshot
        .and_then(provider_upgrade_heartbeat_at)
        .map(|value| value.to_rfc3339_opts(chrono::SecondsFormat::Millis, true));
    let mut object = serde_json::Map::new();
    insert_json_field_with_unavailable_reason(
        &mut object,
        "containerId",
        labels
            .and_then(|labels| provider_upgrade_label_string(labels, "providerGatewayContainerId")),
        unavailable_reason,
    );
    insert_json_field_with_unavailable_reason(
        &mut object,
        "containerName",
        labels.and_then(|labels| {
            provider_upgrade_label_string(labels, "providerGatewayContainerName")
        }),
        unavailable_reason,
    );
    insert_json_field_with_unavailable_reason(
        &mut object,
        "version",
        labels.and_then(|labels| provider_upgrade_label_string(labels, "providerGatewayVersion")),
        unavailable_reason,
    );
    insert_json_field_with_unavailable_reason(
        &mut object,
        "restartPolicy",
        labels.and_then(|labels| {
            provider_upgrade_label_string(labels, "providerGatewayRestartPolicy")
        }),
        unavailable_reason,
    );
    insert_json_field_with_unavailable_reason(
        &mut object,
        "pidMode",
        labels.and_then(|labels| provider_upgrade_label_string(labels, "providerGatewayPidMode")),
        unavailable_reason,
    );
    insert_json_field_with_unavailable_reason(
        &mut object,
        "heartbeatTimestamp",
        heartbeat,
        unavailable_reason,
    );
    Value::Object(object)
}

fn evaluate_provider_upgrade_final_heartbeat(
    expectation: &ProviderUpgradeFinalHeartbeatExpectation,
    snapshot: Option<&ProviderUpgradeNodeSnapshot>,
    now: DateTime<Utc>,
) -> ProviderUpgradeFinalHeartbeatDecision {
    let candidate_snapshot = snapshot.is_some_and(|snapshot| {
        provider_upgrade_snapshot_is_candidate(snapshot, expectation.candidate_name.as_deref())
    });
    let unavailable_reason = if candidate_snapshot {
        "latest heartbeat came from candidate validation container, not final promoted compose service"
    } else {
        "final promoted compose heartbeat not observed"
    };
    let final_container_snapshot = if candidate_snapshot { None } else { snapshot };
    if let Some(snapshot) = snapshot {
        let observed_version =
            provider_upgrade_label_string(&snapshot.labels, "providerGatewayVersion");
        let heartbeat_is_new = provider_upgrade_heartbeat_at(snapshot)
            .is_some_and(|last| last >= expectation.min_heartbeat_at);
        let same_provider = snapshot.provider_id.as_str() == expectation.provider_id.as_str();
        let online = snapshot.status == "online";
        if same_provider
            && online
            && observed_version.as_deref() == Some(expectation.target_version.as_str())
            && heartbeat_is_new
            && !candidate_snapshot
        {
            let restart_policy =
                provider_upgrade_label_string(&snapshot.labels, "providerGatewayRestartPolicy");
            let pid_mode =
                provider_upgrade_label_string(&snapshot.labels, "providerGatewayPidMode");
            if restart_policy.as_deref() != Some("always") || pid_mode.as_deref() != Some("host") {
                return ProviderUpgradeFinalHeartbeatDecision::Failed(json!({
                    "ok": false,
                    "reason": "final-promoted-runtime-guard-failed",
                    "blocker": true,
                    "providerId": expectation.provider_id.as_str(),
                    "taskId": expectation.task_id.as_str(),
                    "targetProviderGatewayVersion": expectation.target_version.as_str(),
                    "finalPromotedContainer": provider_upgrade_final_container_json(Some(snapshot), "runtime guard heartbeat fields are incomplete"),
                    "rollback": {
                        "preferred": "last-known-good",
                        "required": true,
                        "lastKnownGoodProviderGatewayVersion": expectation.last_known_good_version.as_deref(),
                        "updaterRollbackExpected": true
                    },
                    "candidateValidationWasTerminal": false,
                    "schedulerResult": expectation.scheduler_result.clone()
                }));
            }
            return ProviderUpgradeFinalHeartbeatDecision::Succeeded(json!({
                "ok": true,
                "reason": "final-promoted-heartbeat-confirmed",
                "providerId": expectation.provider_id.as_str(),
                "taskId": expectation.task_id.as_str(),
                "targetProviderGatewayVersion": expectation.target_version.as_str(),
                "finalPromotedContainer": provider_upgrade_final_container_json(Some(snapshot), "not reported by provider heartbeat"),
                "heartbeat": {
                    "minHeartbeatAt": expectation.min_heartbeat_at.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
                    "observedAt": provider_upgrade_heartbeat_at(snapshot).map(|value| value.to_rfc3339_opts(chrono::SecondsFormat::Millis, true))
                },
                "candidateValidationWasTerminal": false,
                "schedulerResult": expectation.scheduler_result.clone()
            }));
        }
    }

    let pending_result = json!({
        "ok": false,
        "status": "running",
        "reason": "awaiting-final-promoted-heartbeat",
        "providerId": expectation.provider_id.as_str(),
        "taskId": expectation.task_id.as_str(),
        "targetProviderGatewayVersion": expectation.target_version.as_str(),
        "deadlineAt": expectation.deadline_at.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
        "candidateValidationWasTerminal": false,
        "finalPromotedContainer": provider_upgrade_final_container_json(final_container_snapshot, unavailable_reason),
        "schedulerResult": expectation.scheduler_result.clone()
    });
    if now < expectation.deadline_at {
        return ProviderUpgradeFinalHeartbeatDecision::Pending(pending_result);
    }

    ProviderUpgradeFinalHeartbeatDecision::Failed(json!({
        "ok": false,
        "reason": "final-promoted-heartbeat-timeout",
        "error": "final promoted provider-gateway compose container did not report a target-version heartbeat before the guard window expired",
        "blocker": true,
        "providerId": expectation.provider_id.as_str(),
        "taskId": expectation.task_id.as_str(),
        "targetProviderGatewayVersion": expectation.target_version.as_str(),
        "heartbeat": {
            "minHeartbeatAt": expectation.min_heartbeat_at.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
            "deadlineAt": expectation.deadline_at.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
            "observedAt": snapshot.and_then(provider_upgrade_heartbeat_at).map(|value| value.to_rfc3339_opts(chrono::SecondsFormat::Millis, true))
        },
        "finalPromotedContainer": provider_upgrade_final_container_json(final_container_snapshot, unavailable_reason),
        "rollback": {
            "preferred": "last-known-good",
            "required": true,
            "lastKnownGoodProviderGatewayVersion": expectation.last_known_good_version.as_deref(),
            "updaterRollbackExpected": true
        },
        "candidateValidationWasTerminal": false,
        "schedulerResult": expectation.scheduler_result.clone()
    }))
}

async fn provider_upgrade_node_snapshot(
    state: &Arc<AppState>,
    provider_id: &str,
) -> anyhow::Result<Option<ProviderUpgradeNodeSnapshot>> {
    if !state.db_ready() {
        return Ok(None);
    }
    let client = state.pool.get().await?;
    let row = client
        .query_opt(
            "SELECT provider_id, status, labels FROM unidesk_nodes WHERE provider_id = $1 LIMIT 1",
            &[&provider_id],
        )
        .await?;
    Ok(row.map(|row| ProviderUpgradeNodeSnapshot {
        provider_id: row.get("provider_id"),
        status: row.get("status"),
        labels: row.get("labels"),
    }))
}

async fn complete_provider_upgrade_task(
    state: &Arc<AppState>,
    task_id: &str,
    provider_id: &str,
    status: &str,
    result: Value,
) -> anyhow::Result<()> {
    let client = state.pool.get().await?;
    let updated = client
        .execute(
            r#"
            UPDATE unidesk_tasks
            SET status = $2, result = $3::jsonb, updated_at = now()
            WHERE id = $1
              AND command = 'provider.upgrade'
              AND status NOT IN ('succeeded', 'failed')
            "#,
            &[&task_id, &status, &result],
        )
        .await?;
    record_event(
        state,
        "provider_upgrade_final_heartbeat_result",
        provider_id,
        json!({
            "providerId": provider_id,
            "taskId": task_id,
            "status": status,
            "result": result,
            "updated": updated,
        }),
    )
    .await;
    if updated > 0 && is_terminal_task_status(status) {
        notify_task_terminal(state, task_id).await?;
    }
    Ok(())
}

async fn mark_provider_upgrade_awaiting_final_heartbeat(
    state: &Arc<AppState>,
    expectation: &ProviderUpgradeFinalHeartbeatExpectation,
    pending_result: Value,
) -> anyhow::Result<()> {
    let client = state.pool.get().await?;
    client
        .execute(
            r#"
            UPDATE unidesk_tasks
            SET status = 'running', result = $2::jsonb, updated_at = now()
            WHERE id = $1
              AND command = 'provider.upgrade'
              AND status NOT IN ('succeeded', 'failed')
            "#,
            &[&expectation.task_id, &pending_result],
        )
        .await?;
    record_event(
        state,
        "provider_upgrade_final_heartbeat_wait",
        &expectation.provider_id,
        json!({
            "providerId": expectation.provider_id.as_str(),
            "taskId": expectation.task_id.as_str(),
            "targetProviderGatewayVersion": expectation.target_version.as_str(),
            "deadlineAt": expectation.deadline_at.to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
            "candidateValidationWasTerminal": false,
        }),
    )
    .await;
    Ok(())
}

async fn monitor_provider_upgrade_final_heartbeat(
    state: Arc<AppState>,
    expectation: ProviderUpgradeFinalHeartbeatExpectation,
) -> anyhow::Result<()> {
    loop {
        let snapshot = provider_upgrade_node_snapshot(&state, &expectation.provider_id).await?;
        match evaluate_provider_upgrade_final_heartbeat(&expectation, snapshot.as_ref(), Utc::now())
        {
            ProviderUpgradeFinalHeartbeatDecision::Succeeded(result) => {
                complete_provider_upgrade_task(
                    &state,
                    &expectation.task_id,
                    &expectation.provider_id,
                    "succeeded",
                    result,
                )
                .await?;
                break;
            }
            ProviderUpgradeFinalHeartbeatDecision::Failed(result) => {
                complete_provider_upgrade_task(
                    &state,
                    &expectation.task_id,
                    &expectation.provider_id,
                    "failed",
                    result,
                )
                .await?;
                break;
            }
            ProviderUpgradeFinalHeartbeatDecision::Pending(_) => {
                tokio::time::sleep(std::time::Duration::from_millis(
                    PROVIDER_UPGRADE_FINAL_HEARTBEAT_POLL_MS,
                ))
                .await;
            }
        }
    }
    Ok(())
}

async fn ensure_provider_upgrade_final_heartbeat_watcher(
    state: &Arc<AppState>,
    expectation: ProviderUpgradeFinalHeartbeatExpectation,
) {
    let task_id = expectation.task_id.clone();
    let inserted = {
        let mut active = state.active_provider_upgrade_watchers.lock().await;
        active.insert(task_id.clone())
    };
    if !inserted {
        return;
    }
    let state_for_task = state.clone();
    tokio::spawn(async move {
        let provider_id = expectation.provider_id.clone();
        if let Err(error) =
            monitor_provider_upgrade_final_heartbeat(state_for_task.clone(), expectation).await
        {
            state_for_task.log(
                "error",
                "provider_upgrade_final_heartbeat_watcher_failed",
                Some(json!({ "taskId": task_id.as_str(), "providerId": provider_id.as_str(), "error": error.to_string() })),
            );
        }
        let mut active = state_for_task.active_provider_upgrade_watchers.lock().await;
        active.remove(&task_id);
    });
}

async fn maybe_handle_provider_upgrade_schedule_status(
    state: &Arc<AppState>,
    provider_id: &str,
    task_id: &str,
    status: &str,
    stored_result: &Value,
) -> anyhow::Result<bool> {
    let starts_final_guard = status == "succeeded"
        || (status == "running"
            && stored_result.get("mode").and_then(Value::as_str) == Some("schedule")
            && stored_result
                .get("terminalStatusManagedByBackendCore")
                .and_then(Value::as_bool)
                == Some(true));
    if !starts_final_guard {
        return Ok(false);
    }
    let Some(task) = raw_task(state, task_id).await? else {
        return Ok(false);
    };
    if !provider_upgrade_schedule_task(&task) {
        return Ok(false);
    }
    let now = Utc::now();
    let Some(expectation) = provider_upgrade_expectation(&task, stored_result, now) else {
        let result = json!({
            "ok": false,
            "reason": "target-provider-gateway-version-unavailable",
            "error": "provider.upgrade mode=schedule cannot be finalized because targetProviderGatewayVersion was not reported",
            "blocker": true,
            "providerId": provider_id,
            "taskId": task_id,
            "finalPromotedContainer": provider_upgrade_final_container_json(None, "target gateway version unavailable before final heartbeat guard"),
            "candidateValidationWasTerminal": false,
            "schedulerResult": stored_result
        });
        complete_provider_upgrade_task(state, task_id, provider_id, "failed", result).await?;
        return Ok(true);
    };
    let pending = match evaluate_provider_upgrade_final_heartbeat(&expectation, None, now) {
        ProviderUpgradeFinalHeartbeatDecision::Pending(result) => result,
        ProviderUpgradeFinalHeartbeatDecision::Succeeded(result)
        | ProviderUpgradeFinalHeartbeatDecision::Failed(result) => result,
    };
    mark_provider_upgrade_awaiting_final_heartbeat(state, &expectation, pending).await?;
    ensure_provider_upgrade_final_heartbeat_watcher(state, expectation).await;
    Ok(true)
}

pub async fn provider_ws(
    state: Arc<AppState>,
    ws: WebSocketUpgrade,
    headers: HeaderMap,
    uri: Uri,
) -> Response {
    let query_token = uri.query().and_then(|query| {
        url::form_urlencoded::parse(query.as_bytes())
            .find(|(key, _)| key == "token")
            .map(|(_, value)| value.into_owned())
    });
    let token = query_token.or_else(|| {
        headers
            .get("x-provider-token")
            .and_then(|value| value.to_str().ok())
            .map(ToOwned::to_owned)
    });
    if token.as_deref() != Some(&state.config.provider_token) {
        return json_response(
            json!({ "ok": false, "error": "invalid provider token" }),
            401,
        );
    }
    ws.on_upgrade(move |socket| provider_socket_task(state, socket))
}

async fn provider_socket_task(state: Arc<AppState>, socket: WebSocket) {
    state.log("info", "provider_socket_open", None);
    let (mut sender_ws, mut receiver_ws) = socket.split();
    let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
    let connection = Arc::new(ProviderConnection {
        provider_id: tokio::sync::Mutex::new(None),
        sender: tx.clone(),
    });
    let send_task = tokio::spawn(async move {
        while let Some(message) = rx.recv().await {
            if sender_ws.send(message).await.is_err() {
                break;
            }
        }
    });

    while let Some(message) = receiver_ws.next().await {
        match message {
            Ok(Message::Text(text)) => {
                if let Err(error) = handle_provider_text(&state, &connection, &text).await {
                    state.log(
                        "error",
                        "provider_message_failed",
                        Some(provider_message_error_context(&text, &error)),
                    );
                    let _ = tx.send(Message::Text(json!({ "type": "ack", "requestId": "message", "ok": false, "message": error.to_string() }).to_string()));
                }
            }
            Ok(Message::Binary(bytes)) => {
                if let Ok(text) = String::from_utf8(bytes.to_vec()) {
                    if let Err(error) = handle_provider_text(&state, &connection, &text).await {
                        state.log(
                            "error",
                            "provider_message_failed",
                            Some(provider_message_error_context(&text, &error)),
                        );
                    }
                }
            }
            Ok(Message::Close(_)) | Err(_) => break,
            _ => {}
        }
    }

    let provider_id = connection.provider_id.lock().await.clone();
    state.log(
        "warn",
        "provider_socket_close",
        Some(json!({ "providerId": provider_id })),
    );
    if let Some(provider_id) = provider_id {
        let should_mark_offline = {
            let mut providers = state.active_providers.write().await;
            if providers
                .get(&provider_id)
                .is_some_and(|active| Arc::ptr_eq(active, &connection))
            {
                providers.remove(&provider_id);
                true
            } else {
                false
            }
        };
        crate::egress_tcp::close_egress_tcp_connections_for_provider(&state, &provider_id).await;
        if should_mark_offline {
            if let Err(error) = mark_provider_offline(&state, &provider_id).await {
                state.log(
                    "error",
                    "provider_offline_mark_failed",
                    Some(json!({ "providerId": provider_id, "error": error.to_string() })),
                );
            }
        }
    }
    send_task.abort();
}

fn provider_message_error_context(text: &str, error: &anyhow::Error) -> Value {
    let parsed = serde_json::from_str::<Value>(text).ok();
    json!({
        "error": error.to_string(),
        "messageType": parsed
            .as_ref()
            .and_then(|value| value.get("type"))
            .and_then(Value::as_str)
            .unwrap_or("unknown"),
        "providerId": parsed
            .as_ref()
            .and_then(|value| value.get("providerId"))
            .and_then(Value::as_str)
            .unwrap_or("unknown"),
        "requestId": parsed
            .as_ref()
            .and_then(|value| value.get("requestId"))
            .and_then(Value::as_str)
            .unwrap_or("unknown"),
        "textBytes": text.len(),
        "textPreview": text.chars().take(300).collect::<String>(),
    })
}

async fn handle_provider_text(
    state: &Arc<AppState>,
    connection: &Arc<ProviderConnection>,
    text: &str,
) -> anyhow::Result<()> {
    let message: Value = serde_json::from_str(text)?;
    let provider_id = message
        .get("providerId")
        .and_then(Value::as_str)
        .unwrap_or("")
        .to_string();
    if provider_id.is_empty() {
        anyhow::bail!(
            "Unsupported provider message: {}",
            &text[..text.len().min(200)]
        );
    }
    {
        let mut id = connection.provider_id.lock().await;
        *id = Some(provider_id.clone());
    }
    state
        .active_providers
        .write()
        .await
        .insert(provider_id.clone(), connection.clone());
    let message_type = message.get("type").and_then(Value::as_str).unwrap_or("");
    match message_type {
        "host_ssh_error" => {
            forward_ssh_provider_message(state, &message).await;
        }
        "http_tunnel_response" => {
            let request_id = message
                .get("requestId")
                .and_then(Value::as_str)
                .unwrap_or("")
                .to_string();
            let waiter = state.http_tunnel_waiters.lock().await.remove(&request_id);
            if let Some(waiter) = waiter {
                let _ = waiter.send(Ok(HttpTunnelResponse {
                    provider_id,
                    request_id,
                    ok: message.get("ok").and_then(Value::as_bool).unwrap_or(false),
                    result: message.get("result").cloned().unwrap_or(Value::Null),
                }));
            } else {
                state.log(
                    "warn",
                    "http_tunnel_response_without_waiter",
                    Some(json!({ "requestId": request_id })),
                );
            }
        }
        "egress_tcp_open" => handle_egress_tcp_open(state, connection.clone(), &message).await?,
        "egress_tcp_data" => handle_egress_tcp_data(state, &message).await?,
        "egress_tcp_close" => handle_egress_tcp_close(state, &message).await?,
        "register" => {
            let mut labels = message.get("labels").cloned().unwrap_or_else(|| json!({}));
            if let Some(object) = labels.as_object_mut() {
                object.insert(
                    "unideskCapabilities".to_string(),
                    message
                        .get("capabilities")
                        .cloned()
                        .unwrap_or_else(|| json!([])),
                );
            }
            let name = message
                .get("name")
                .and_then(Value::as_str)
                .unwrap_or(&provider_id);
            upsert_provider_node(state, &provider_id, name, &labels).await?;
            record_event(state, "provider_registered", &provider_id, json!({
                "providerId": provider_id,
                "name": name,
                "labels": labels,
                "capabilities": message.get("capabilities").cloned().unwrap_or_else(|| json!([])),
            })).await;
            let _ = connection.sender.send(Message::Text(json!({ "type": "ack", "requestId": "register", "ok": true, "message": "registered" }).to_string()));
        }
        "heartbeat" => {
            let mut labels = message.get("labels").cloned().unwrap_or_else(|| json!({}));
            if let Some(object) = labels.as_object_mut() {
                object.insert(
                    "providerGatewayHeartbeatAt".to_string(),
                    message
                        .get("at")
                        .cloned()
                        .unwrap_or_else(|| Value::String(Utc::now().to_rfc3339())),
                );
            }
            update_provider_heartbeat(state, &provider_id, &labels).await?;
            state.log(
                "debug",
                "provider_heartbeat",
                Some(json!({ "providerId": provider_id, "labels": labels })),
            );
        }
        "system_status" => {
            let status = message.get("status").cloned().unwrap_or_else(|| json!({}));
            let collected_at = status
                .get("collectedAt")
                .and_then(Value::as_str)
                .unwrap_or_else(|| message.get("at").and_then(Value::as_str).unwrap_or(""));
            upsert_system_status(state, &provider_id, &status, collected_at).await?;
            state.log("debug", "provider_system_status", Some(json!({
                "providerId": provider_id,
                "cpuPercent": status.pointer("/cpu/percent").cloned().unwrap_or(Value::Null),
                "memoryPercent": status.pointer("/memory/percent").cloned().unwrap_or(Value::Null),
                "diskPercent": status.pointer("/disk/percent").cloned().unwrap_or(Value::Null),
                "ok": status.get("ok").cloned().unwrap_or(Value::Null),
            })));
        }
        "docker_status" => {
            let status = message.get("status").cloned().unwrap_or_else(|| json!({}));
            let collected_at = status
                .get("collectedAt")
                .and_then(Value::as_str)
                .unwrap_or_else(|| message.get("at").and_then(Value::as_str).unwrap_or(""));
            upsert_docker_status(state, &provider_id, &status, collected_at).await?;
            state.log("debug", "provider_docker_status", Some(json!({ "providerId": provider_id, "counts": status.get("counts").cloned().unwrap_or(Value::Null), "ok": status.get("ok").cloned().unwrap_or(Value::Null) })));
        }
        "task_status" => {
            let task_id = message.get("taskId").and_then(Value::as_str).unwrap_or("");
            let status = message.get("status").and_then(Value::as_str).unwrap_or("");
            let stored_result =
                compact_json_for_storage(&message.get("result").cloned().unwrap_or_else(
                    || json!({ "message": message.get("message").cloned().unwrap_or(Value::Null) }),
                ));
            if maybe_handle_provider_upgrade_schedule_status(
                state,
                &provider_id,
                task_id,
                status,
                &stored_result,
            )
            .await?
            {
                return Ok(());
            }
            let client = state.pool.get().await?;
            client.execute(
                r#"
                WITH incoming AS (
                  SELECT $2::text AS status, $3::jsonb AS result
                )
                UPDATE unidesk_tasks
                SET
                  status = CASE
                    WHEN unidesk_tasks.status IN ('succeeded', 'failed') AND incoming.status NOT IN ('succeeded', 'failed') THEN unidesk_tasks.status
                    WHEN unidesk_tasks.status = 'running' AND incoming.status = 'accepted' THEN unidesk_tasks.status
                    ELSE incoming.status
                  END,
                  result = CASE
                    WHEN unidesk_tasks.status IN ('succeeded', 'failed') AND incoming.status NOT IN ('succeeded', 'failed') THEN unidesk_tasks.result
                    WHEN unidesk_tasks.status = 'running' AND incoming.status = 'accepted' THEN unidesk_tasks.result
                    ELSE incoming.result
                  END,
                  updated_at = now()
                FROM incoming
                WHERE id = $1
                "#,
                &[&task_id, &status, &stored_result],
            ).await?;
            record_event(
                state,
                "task_status",
                &provider_id,
                json!({
                    "providerId": provider_id,
                    "taskId": task_id,
                    "status": status,
                    "message": message.get("message").cloned().unwrap_or(Value::Null),
                    "result": stored_result,
                }),
            )
            .await;
            if is_terminal_task_status(status) {
                notify_task_terminal(state, task_id).await?;
            }
        }
        _ => anyhow::bail!(
            "Unsupported provider message: {}",
            &text[..text.len().min(200)]
        ),
    }
    Ok(())
}

pub async fn mark_provider_offline(state: &Arc<AppState>, provider_id: &str) -> anyhow::Result<()> {
    state.active_providers.write().await.remove(provider_id);
    state
        .active_ssh_data_channels
        .lock()
        .await
        .retain(|_, channel| channel.provider_id != provider_id);
    if !state.db_ready() {
        return Ok(());
    }
    let client = state.pool.get().await?;
    client.execute("UPDATE unidesk_nodes SET status = 'offline', updated_at = now() WHERE provider_id = $1", &[&provider_id]).await?;
    record_event(
        state,
        "provider_offline",
        provider_id,
        json!({ "providerId": provider_id }),
    )
    .await;
    Ok(())
}

pub async fn mark_stale_providers_offline(state: &Arc<AppState>) -> anyhow::Result<()> {
    if !state.db_ready() {
        return Ok(());
    }
    let timeout_ms = state.config.heartbeat_timeout_ms as i64;
    let cutoff = Utc::now() - Duration::milliseconds(timeout_ms);
    let client = state.pool.get().await?;
    let rows = client
        .query(
            r#"
        UPDATE unidesk_nodes
        SET status = 'offline', updated_at = now()
        WHERE status = 'online'
          AND last_heartbeat IS NOT NULL
          AND last_heartbeat < $1
        RETURNING provider_id
        "#,
            &[&cutoff],
        )
        .await?;
    for row in rows {
        let provider_id: String = row.get("provider_id");
        state.active_providers.write().await.remove(&provider_id);
        record_event(
            state,
            "provider_heartbeat_timeout",
            &provider_id,
            json!({ "providerId": provider_id, "timeoutMs": timeout_ms }),
        )
        .await;
    }
    Ok(())
}

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

    fn dt(value: &str) -> DateTime<Utc> {
        DateTime::parse_from_rfc3339(value)
            .expect("valid fixture timestamp")
            .with_timezone(&Utc)
    }

    fn expectation() -> ProviderUpgradeFinalHeartbeatExpectation {
        ProviderUpgradeFinalHeartbeatExpectation {
            task_id: "task_upgrade_contract".to_string(),
            provider_id: "D601".to_string(),
            target_version: "0.2.27".to_string(),
            min_heartbeat_at: dt("2026-05-24T10:00:00Z"),
            deadline_at: dt("2026-05-24T10:04:00Z"),
            scheduler_result: json!({
                "mode": "schedule",
                "providerGatewayVersion": "0.2.26",
                "targetProviderGatewayVersion": "0.2.27",
                "plan": {
                    "candidateName": "unidesk-provider-gateway-D601-candidate-task-upgrade-contract"
                }
            }),
            candidate_name: Some(
                "unidesk-provider-gateway-D601-candidate-task-upgrade-contract".to_string(),
            ),
            last_known_good_version: Some("0.2.26".to_string()),
        }
    }

    #[test]
    fn candidate_validation_heartbeat_is_not_terminal_success() {
        let snapshot = ProviderUpgradeNodeSnapshot {
            provider_id: "D601".to_string(),
            status: "online".to_string(),
            labels: json!({
                "providerGatewayVersion": "0.2.27",
                "providerGatewayContainerId": "candidate-container-id",
                "providerGatewayContainerName": "unidesk-provider-gateway-D601-candidate-task-upgrade-contract",
                "providerGatewayRestartPolicy": "no",
                "providerGatewayPidMode": "host",
                "providerGatewayHeartbeatAt": "2026-05-24T10:01:00Z"
            }),
        };
        let decision = evaluate_provider_upgrade_final_heartbeat(
            &expectation(),
            Some(&snapshot),
            dt("2026-05-24T10:05:00Z"),
        );
        let ProviderUpgradeFinalHeartbeatDecision::Failed(result) = decision else {
            panic!("candidate heartbeat must fail after final promoted heartbeat timeout");
        };
        assert_eq!(
            result.get("reason").and_then(Value::as_str),
            Some("final-promoted-heartbeat-timeout")
        );
        assert_eq!(
            result
                .get("candidateValidationWasTerminal")
                .and_then(Value::as_bool),
            Some(false)
        );
        assert_eq!(
            result
                .pointer("/rollback/preferred")
                .and_then(Value::as_str),
            Some("last-known-good")
        );
        assert_eq!(
            result
                .pointer("/finalPromotedContainer/containerIdUnavailableReason")
                .and_then(Value::as_str),
            Some("latest heartbeat came from candidate validation container, not final promoted compose service")
        );
    }

    #[test]
    fn final_promoted_compose_heartbeat_is_success() {
        let snapshot = ProviderUpgradeNodeSnapshot {
            provider_id: "D601".to_string(),
            status: "online".to_string(),
            labels: json!({
                "providerGatewayVersion": "0.2.27",
                "providerGatewayContainerId": "final-container-id",
                "providerGatewayContainerName": "unidesk-provider-gateway-D601",
                "providerGatewayRestartPolicy": "always",
                "providerGatewayPidMode": "host",
                "providerGatewayHeartbeatAt": "2026-05-24T10:02:00Z"
            }),
        };
        let decision = evaluate_provider_upgrade_final_heartbeat(
            &expectation(),
            Some(&snapshot),
            dt("2026-05-24T10:02:01Z"),
        );
        let ProviderUpgradeFinalHeartbeatDecision::Succeeded(result) = decision else {
            panic!("final compose heartbeat should succeed");
        };
        assert_eq!(
            result.get("reason").and_then(Value::as_str),
            Some("final-promoted-heartbeat-confirmed")
        );
        assert_eq!(
            result
                .pointer("/finalPromotedContainer/containerId")
                .and_then(Value::as_str),
            Some("final-container-id")
        );
        assert_eq!(
            result
                .pointer("/finalPromotedContainer/heartbeatTimestamp")
                .and_then(Value::as_str),
            Some("2026-05-24T10:02:00.000Z")
        );
    }
}
