use std::sync::Arc;

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::http::{HeaderMap, Uri};
use axum::response::Response;
use base64::Engine;
use futures_util::{SinkExt, StreamExt};
use rand::Rng;
use serde_json::{json, Value};
use tokio::sync::mpsc;

use crate::db::provider_supports;
use crate::http::json_response;
use crate::ssh_data_channel::{
    active_ssh_data_channel_count, claim_ssh_data_channel, close_ssh_data_session,
    idle_ssh_data_channel_count, remove_ssh_data_channel, send_ssh_data_frame,
};
use crate::state::{AppState, SshClientConnection, SshDataFrame};

pub async fn ssh_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 ssh bridge token" }),
            401,
        );
    }
    ws.on_upgrade(move |socket| ssh_socket_task(state, socket))
}

async fn ssh_socket_task(state: Arc<AppState>, socket: WebSocket) {
    let (mut sender_ws, mut receiver_ws) = socket.split();
    let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
    let send_task = tokio::spawn(async move {
        while let Some(message) = rx.recv().await {
            if sender_ws.send(message).await.is_err() {
                break;
            }
        }
    });
    let mut session_id: Option<String> = None;
    while let Some(message) = receiver_ws.next().await {
        match message {
            Ok(Message::Text(text)) => {
                if let Err(error) =
                    handle_ssh_client_text(&state, tx.clone(), &mut session_id, text).await
                {
                    let _ = tx.send(Message::Text(
                        json!({ "type": "ssh.error", "message": error.to_string() }).to_string(),
                    ));
                }
            }
            Ok(Message::Close(_)) | Err(_) => break,
            _ => {}
        }
    }
    if let Some(session_id) = session_id {
        close_ssh_client(&state, &session_id, Some("ssh client disconnected")).await;
    }
    send_task.abort();
}

async fn handle_ssh_client_text(
    state: &Arc<AppState>,
    client_sender: mpsc::UnboundedSender<Message>,
    session_id: &mut Option<String>,
    text: String,
) -> anyhow::Result<()> {
    let message: Value = serde_json::from_str(&text)?;
    let message_type = message.get("type").and_then(Value::as_str).unwrap_or("");
    if message_type == "ssh.open" {
        let provider_id = message
            .get("providerId")
            .and_then(Value::as_str)
            .unwrap_or("");
        if provider_id.is_empty() {
            let _ = client_sender.send(Message::Text(
                json!({ "type": "ssh.error", "message": "providerId is required" }).to_string(),
            ));
            return Ok(());
        }
        let provider = {
            let providers = state.active_providers.read().await;
            providers.get(provider_id).cloned()
        };
        let Some(provider) = provider else {
            let _ = client_sender.send(Message::Text(json!({ "type": "ssh.error", "message": format!("provider is not online: {provider_id}") }).to_string()));
            return Ok(());
        };
        if !provider_supports(state, provider_id, "host.ssh")
            .await
            .unwrap_or(false)
        {
            let _ = client_sender.send(Message::Text(json!({ "type": "ssh.error", "message": format!("provider does not declare host.ssh capability: {provider_id}") }).to_string()));
            return Ok(());
        }
        if !provider_supports(state, provider_id, "host.ssh.tcp-pool")
            .await
            .unwrap_or(false)
        {
            let _ = client_sender.send(Message::Text(json!({
                "type": "ssh.error",
                "message": format!("provider gateway must be upgraded before host.ssh can run: {provider_id} does not declare host.ssh.tcp-pool"),
                "failureKind": "provider-gateway-upgrade-required",
                "providerId": provider_id,
                "requiredCapability": "host.ssh.tcp-pool",
            }).to_string()));
            return Ok(());
        }
        let new_session_id = safe_session_id();
        let Some((data_channel_id, _data_writer)) =
            claim_ssh_data_channel(state, provider_id, &new_session_id).await
        else {
            let ready = idle_ssh_data_channel_count(state, provider_id).await;
            let total = active_ssh_data_channel_count(state, provider_id).await;
            let _ = client_sender.send(Message::Text(json!({
                "type": "ssh.error",
                "message": format!("provider ssh tcp data pool has no idle channel: {provider_id}"),
                "failureKind": "provider-data-pool-exhausted",
                "providerId": provider_id,
                "dataPool": {
                    "ready": ready,
                    "total": total,
                    "required": 1
                },
                "requiredCapability": "host.ssh.tcp-pool",
            }).to_string()));
            return Ok(());
        };
        *session_id = Some(new_session_id.clone());
        state.active_ssh_clients.write().await.insert(
            new_session_id.clone(),
            SshClientConnection {
                provider_id: provider_id.to_string(),
                sender: client_sender.clone(),
                data_channel_id: data_channel_id.clone(),
            },
        );
        let mut open_message = json!({
            "type": "host_ssh_open",
            "sessionId": new_session_id,
            "transport": "tcp-pool",
            "dataChannelId": data_channel_id,
            "cols": number_from_unknown(message.get("cols"), 100, 20, 300),
            "rows": number_from_unknown(message.get("rows"), 30, 8, 120),
        });
        if let Some(cwd) = message
            .get("cwd")
            .and_then(Value::as_str)
            .filter(|value| !value.is_empty())
        {
            open_message["cwd"] = Value::String(cwd.to_string());
        }
        if let Some(command) = message
            .get("command")
            .and_then(Value::as_str)
            .filter(|value| !value.is_empty())
        {
            open_message["command"] = Value::String(command.to_string());
        }
        if let Some(tty) = message.get("tty").and_then(Value::as_bool) {
            open_message["tty"] = Value::Bool(tty);
        }
        let _ = provider
            .sender
            .send(Message::Text(open_message.to_string()));
        let _ = client_sender.send(Message::Text(json!({ "type": "ssh.dispatched", "providerId": provider_id, "sessionId": new_session_id, "transport": "tcp-pool", "dataChannelId": open_message.get("dataChannelId").cloned().unwrap_or(Value::Null) }).to_string()));
        state.log("info", "ssh_session_dispatched", Some(json!({ "providerId": provider_id, "sessionId": new_session_id, "transport": "tcp-pool", "dataChannelId": open_message.get("dataChannelId").cloned().unwrap_or(Value::Null), "hasCommand": open_message.get("command").is_some() })));
        return Ok(());
    }

    let Some(current_session_id) = session_id.clone() else {
        if message_type == "ssh.eof" || message_type == "ssh.close" {
            return Ok(());
        }
        let _ = client_sender.send(Message::Text(
            json!({ "type": "ssh.error", "message": "ssh session is not open" }).to_string(),
        ));
        return Ok(());
    };
    let session = {
        let clients = state.active_ssh_clients.read().await;
        clients
            .get(&current_session_id)
            .map(|client| (client.provider_id.clone(), client.data_channel_id.clone()))
    };
    let Some((provider_id, data_channel_id)) = session else {
        let _ = client_sender.send(Message::Text(
            json!({ "type": "ssh.error", "message": "ssh session disappeared" }).to_string(),
        ));
        return Ok(());
    };

    match message_type {
        "ssh.input" => {
            let data = message.get("data").and_then(Value::as_str).unwrap_or("");
            if message.get("encoding").and_then(Value::as_str) != Some("base64") || data.is_empty()
            {
                let _ = client_sender.send(Message::Text(
                    json!({ "type": "ssh.error", "message": "ssh.input requires base64 data" })
                        .to_string(),
                ));
                return Ok(());
            }
            let bytes = match base64::engine::general_purpose::STANDARD.decode(data) {
                Ok(bytes) => bytes,
                Err(error) => {
                    let _ = client_sender.send(Message::Text(
                        json!({ "type": "ssh.error", "message": format!("invalid ssh.input base64: {error}") })
                            .to_string(),
                    ));
                    return Ok(());
                }
            };
            if !send_ssh_data_frame(
                state,
                &provider_id,
                &data_channel_id,
                SshDataFrame {
                    header: json!({ "type": "input", "sessionId": current_session_id }),
                    payload: bytes,
                },
            )
            .await
            {
                let _ = client_sender.send(Message::Text(json!({ "type": "ssh.error", "message": "ssh tcp data channel is not available", "failureKind": "provider-data-channel-missing", "providerId": provider_id, "dataChannelId": data_channel_id }).to_string()));
            }
        }
        "ssh.resize" => {
            let _ = send_ssh_data_frame(
                state,
                &provider_id,
                &data_channel_id,
                SshDataFrame {
                    header: json!({
                    "type": "resize",
                    "sessionId": current_session_id,
                    "cols": number_from_unknown(message.get("cols"), 100, 20, 300),
                    "rows": number_from_unknown(message.get("rows"), 30, 8, 120),
                    }),
                    payload: Vec::new(),
                },
            )
            .await;
        }
        "ssh.eof" => {
            let _ = send_ssh_data_frame(
                state,
                &provider_id,
                &data_channel_id,
                SshDataFrame {
                    header: json!({ "type": "eof", "sessionId": current_session_id }),
                    payload: Vec::new(),
                },
            )
            .await;
        }
        "ssh.close" => {
            close_ssh_client(state, &current_session_id, None).await;
        }
        _ => {
            let _ = client_sender.send(Message::Text(json!({ "type": "ssh.error", "message": format!("unsupported ssh client message: {message_type}") }).to_string()));
        }
    }
    Ok(())
}

pub async fn forward_ssh_provider_message(state: &Arc<AppState>, message: &Value) {
    let session_id = message
        .get("sessionId")
        .and_then(Value::as_str)
        .unwrap_or("");
    let message_type = message.get("type").and_then(Value::as_str).unwrap_or("");
    let message_text = message
        .get("message")
        .and_then(Value::as_str)
        .unwrap_or("ssh session error");
    let client = {
        let clients = state.active_ssh_clients.read().await;
        clients.get(session_id).map(|client| {
            (
                client.provider_id.clone(),
                client.data_channel_id.clone(),
                client.sender.clone(),
            )
        })
    };
    let Some((provider_id, data_channel_id, client_sender)) = client else {
        state.log("warn", "ssh_client_missing", Some(json!({ "providerId": message.get("providerId").cloned().unwrap_or(Value::Null), "sessionId": session_id, "type": message.get("type").cloned().unwrap_or(Value::Null) })));
        return;
    };
    let stale_channel_not_ready = message_type == "host_ssh_error"
        && message_text.contains("requested ssh tcp data channel is not ready");
    if message_type == "host_ssh_error" {
        let removed = remove_ssh_data_channel(state, &provider_id, &data_channel_id).await;
        state.log(
            if removed { "warn" } else { "info" },
            "ssh_data_channel_removed_after_control_error",
            Some(json!({
                "providerId": provider_id.clone(),
                "sessionId": session_id,
                "dataChannelId": data_channel_id.clone(),
                "removed": removed,
                "message": message_text,
            })),
        );
    }
    let mut outbound = json!({
        "type": "ssh.error",
        "message": message_text,
        "transport": "tcp-pool",
        "controlFallback": true,
    });
    if stale_channel_not_ready {
        outbound["failureKind"] = Value::String("provider-data-channel-missing".to_string());
        outbound["providerId"] = Value::String(provider_id);
        outbound["dataChannelId"] = Value::String(data_channel_id);
    }
    let _ = client_sender.send(Message::Text(outbound.to_string()));
    if message_type == "host_ssh_error" {
        let state = state.clone();
        let session_id = session_id.to_string();
        tokio::spawn(async move {
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
            close_ssh_client(&state, &session_id, None).await;
        });
    }
}

async fn close_ssh_client(state: &Arc<AppState>, session_id: &str, message: Option<&str>) {
    close_ssh_data_session(state, session_id).await;
    let client = state.active_ssh_clients.write().await.remove(session_id);
    if let Some(client) = client {
        let close = Message::Close(Some(axum::extract::ws::CloseFrame {
            code: axum::extract::ws::close_code::NORMAL,
            reason: message.unwrap_or("ssh session closed").to_string().into(),
        }));
        let _ = client.sender.send(close);
    }
}

fn safe_session_id() -> String {
    let suffix: u64 = rand::thread_rng().gen();
    format!(
        "ssh_{}_{}",
        chrono::Utc::now().timestamp_millis(),
        format!("{suffix:x}")
    )
}

fn number_from_unknown(value: Option<&Value>, fallback: i64, min: i64, max: i64) -> i64 {
    let parsed = value
        .and_then(|value| {
            value
                .as_i64()
                .or_else(|| value.as_str().and_then(|text| text.parse::<i64>().ok()))
        })
        .unwrap_or(fallback);
    parsed.clamp(min, max)
}
