use std::io::{Error, ErrorKind};
use std::sync::Arc;

use axum::extract::ws::Message;
use base64::Engine;
use serde_json::{json, Value};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{tcp::OwnedReadHalf, tcp::OwnedWriteHalf, TcpListener, TcpStream};
use tokio::sync::mpsc;

use crate::state::{AppState, SshDataChannel, SshDataFrame};

const SSH_DATA_PROTOCOL: &str = "unidesk-host-ssh-tcp-pool-v1";
const SSH_DATA_MAX_HEADER_BYTES: usize = 64 * 1024;
const SSH_DATA_MAX_PAYLOAD_BYTES: usize = 16 * 1024 * 1024;

fn channel_key(provider_id: &str, channel_id: &str) -> String {
    format!("{provider_id}:{channel_id}")
}

fn header_string(header: &Value, key: &str) -> String {
    header
        .get(key)
        .and_then(Value::as_str)
        .unwrap_or("")
        .to_string()
}

fn release_channel_for_session(channel: &mut SshDataChannel, session_id: &str) -> bool {
    if channel.active_session_id.as_deref() != Some(session_id) {
        return false;
    }
    channel.active_session_id = None;
    true
}

async fn read_data_frame(reader: &mut OwnedReadHalf) -> std::io::Result<Option<(Value, Vec<u8>)>> {
    let mut header_len_bytes = [0_u8; 4];
    if let Err(error) = reader.read_exact(&mut header_len_bytes).await {
        if error.kind() == ErrorKind::UnexpectedEof {
            return Ok(None);
        }
        return Err(error);
    }
    let header_len = u32::from_be_bytes(header_len_bytes) as usize;
    if header_len == 0 || header_len > SSH_DATA_MAX_HEADER_BYTES {
        return Err(Error::new(
            ErrorKind::InvalidData,
            format!("invalid ssh data header length: {header_len}"),
        ));
    }
    let mut header_bytes = vec![0_u8; header_len];
    reader.read_exact(&mut header_bytes).await?;
    let header: Value = serde_json::from_slice(&header_bytes)
        .map_err(|error| Error::new(ErrorKind::InvalidData, error.to_string()))?;
    let payload_len = header
        .get("length")
        .and_then(Value::as_u64)
        .unwrap_or(0) as usize;
    if payload_len > SSH_DATA_MAX_PAYLOAD_BYTES {
        return Err(Error::new(
            ErrorKind::InvalidData,
            format!("ssh data payload too large: {payload_len}"),
        ));
    }
    let mut payload = vec![0_u8; payload_len];
    if payload_len > 0 {
        reader.read_exact(&mut payload).await?;
    }
    Ok(Some((header, payload)))
}

async fn write_data_frame(
    writer: &mut OwnedWriteHalf,
    frame: SshDataFrame,
) -> std::io::Result<()> {
    let mut header = frame.header;
    if let Some(object) = header.as_object_mut() {
        object.insert(
            "length".to_string(),
            Value::Number(serde_json::Number::from(frame.payload.len())),
        );
    }
    let header_bytes = serde_json::to_vec(&header)
        .map_err(|error| Error::new(ErrorKind::InvalidData, error.to_string()))?;
    if header_bytes.len() > SSH_DATA_MAX_HEADER_BYTES {
        return Err(Error::new(
            ErrorKind::InvalidData,
            format!("ssh data header too large: {}", header_bytes.len()),
        ));
    }
    writer
        .write_all(&(header_bytes.len() as u32).to_be_bytes())
        .await?;
    writer.write_all(&header_bytes).await?;
    if !frame.payload.is_empty() {
        writer.write_all(&frame.payload).await?;
    }
    writer.flush().await?;
    Ok(())
}

pub async fn claim_ssh_data_channel(
    state: &Arc<AppState>,
    provider_id: &str,
    session_id: &str,
) -> Option<(String, mpsc::UnboundedSender<SshDataFrame>)> {
    let mut channels = state.active_ssh_data_channels.lock().await;
    let channel = channels
        .values_mut()
        .filter(|channel| channel.provider_id == provider_id && channel.active_session_id.is_none())
        // Prefer the newest channel so a freshly reconnected pool is not
        // starved behind older stale entries that have not failed out yet.
        .max_by_key(|channel| channel.connected_at_millis)?;
    channel.active_session_id = Some(session_id.to_string());
    Some((channel.channel_id.clone(), channel.writer.clone()))
}

pub async fn send_ssh_data_frame(
    state: &Arc<AppState>,
    provider_id: &str,
    channel_id: &str,
    frame: SshDataFrame,
) -> bool {
    let key = channel_key(provider_id, channel_id);
    let mut channels = state.active_ssh_data_channels.lock().await;
    let Some(channel) = channels.get(&key) else {
        return false;
    };
    if channel.writer.send(frame).is_ok() {
        return true;
    }
    channels.remove(&key);
    false
}

pub async fn release_ssh_data_channel(
    state: &Arc<AppState>,
    provider_id: &str,
    channel_id: &str,
    session_id: &str,
) -> bool {
    let mut channels = state.active_ssh_data_channels.lock().await;
    if let Some(channel) = channels.get_mut(&channel_key(provider_id, channel_id)) {
        return release_channel_for_session(channel, session_id);
    }
    false
}

pub async fn remove_ssh_data_channel(
    state: &Arc<AppState>,
    provider_id: &str,
    channel_id: &str,
) -> bool {
    state
        .active_ssh_data_channels
        .lock()
        .await
        .remove(&channel_key(provider_id, channel_id))
        .is_some()
}

pub async fn close_ssh_data_session(state: &Arc<AppState>, session_id: &str) {
    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(),
            )
        })
    };
    let Some((provider_id, channel_id)) = client else {
        return;
    };
    let key = channel_key(&provider_id, &channel_id);
    let mut channels = state.active_ssh_data_channels.lock().await;
    let Some(channel) = channels.get_mut(&key) else {
        return;
    };
    if channel.active_session_id.as_deref() != Some(session_id) {
        state.log(
            "info",
            "ssh_data_channel_close_ignored_stale_session",
            Some(json!({
                "providerId": provider_id,
                "channelId": channel_id,
                "sessionId": session_id,
                "activeSessionId": channel.active_session_id.clone(),
            })),
        );
        return;
    }
    let close = SshDataFrame {
        header: json!({ "type": "close", "sessionId": session_id }),
        payload: Vec::new(),
    };
    if channel.writer.send(close).is_err() {
        channels.remove(&key);
        return;
    }
    let released = release_channel_for_session(channel, session_id);
    debug_assert!(released);
}

pub async fn active_ssh_data_channel_count(state: &Arc<AppState>, provider_id: &str) -> usize {
    let channels = state.active_ssh_data_channels.lock().await;
    channels
        .values()
        .filter(|channel| channel.provider_id == provider_id)
        .count()
}

pub async fn idle_ssh_data_channel_count(state: &Arc<AppState>, provider_id: &str) -> usize {
    let channels = state.active_ssh_data_channels.lock().await;
    channels
        .values()
        .filter(|channel| channel.provider_id == provider_id && channel.active_session_id.is_none())
        .count()
}

async fn forward_provider_frame(
    state: &Arc<AppState>,
    provider_id: &str,
    channel_id: &str,
    header: Value,
    payload: Vec<u8>,
) {
    let session_id = header_string(&header, "sessionId");
    if session_id.is_empty() {
        state.log(
            "warn",
            "ssh_data_frame_missing_session",
            Some(json!({ "providerId": provider_id, "channelId": channel_id, "header": header })),
        );
        return;
    }
    let clients = state.active_ssh_clients.read().await;
    let Some(client) = clients.get(&session_id) else {
        state.log(
            "warn",
            "ssh_data_client_missing",
            Some(json!({ "providerId": provider_id, "channelId": channel_id, "sessionId": session_id, "header": header })),
        );
        return;
    };
    let frame_type = header_string(&header, "type");
    let client_sender = client.sender.clone();
    let outbound = match frame_type.as_str() {
        "opened" => json!({
            "type": "ssh.opened",
            "providerId": provider_id,
            "sessionId": session_id,
            "transport": "tcp-pool",
            "dataChannelId": channel_id,
        }),
        "data" => json!({
            "type": "ssh.data",
            "stream": header.get("stream").cloned().unwrap_or_else(|| Value::String("stdout".to_string())),
            "data": base64::engine::general_purpose::STANDARD.encode(payload),
            "encoding": "base64",
            "transport": "tcp-pool",
            "dataChannelId": channel_id,
        }),
        "exit" => json!({
            "type": "ssh.exit",
            "exitCode": header.get("exitCode").cloned().unwrap_or(Value::Null),
            "signal": header.get("signal").cloned().unwrap_or(Value::Null),
            "transport": "tcp-pool",
            "dataChannelId": channel_id,
        }),
        "error" => json!({
            "type": "ssh.error",
            "message": header.get("message").and_then(Value::as_str).unwrap_or("ssh data channel error"),
            "transport": "tcp-pool",
            "dataChannelId": channel_id,
        }),
        _ => json!({
            "type": "ssh.error",
            "message": format!("unsupported ssh data frame: {frame_type}"),
            "transport": "tcp-pool",
            "dataChannelId": channel_id,
        }),
    };
    let _ = client_sender.send(Message::Text(outbound.to_string()));
    drop(clients);
    if frame_type == "exit" || frame_type == "error" {
        let close = Message::Close(Some(axum::extract::ws::CloseFrame {
            code: axum::extract::ws::close_code::NORMAL,
            reason: "ssh session closed".to_string().into(),
        }));
        let _ = client_sender.send(close);
        if !release_ssh_data_channel(state, provider_id, channel_id, &session_id).await {
            state.log(
                "info",
                "ssh_data_channel_release_ignored_stale_session",
                Some(json!({
                    "providerId": provider_id,
                    "channelId": channel_id,
                    "sessionId": session_id,
                })),
            );
        }
        let state = state.clone();
        tokio::spawn(async move {
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
            state.active_ssh_clients.write().await.remove(&session_id);
        });
    }
}

async fn handle_ssh_data_stream(state: Arc<AppState>, stream: TcpStream) {
    let peer_addr = stream.peer_addr().ok();
    let (mut reader, mut writer) = stream.into_split();
    let hello = match read_data_frame(&mut reader).await {
        Ok(Some((header, _))) => header,
        Ok(None) => return,
        Err(error) => {
            state.log("warn", "ssh_data_hello_read_failed", Some(json!({ "peer": peer_addr.map(|addr| addr.to_string()), "error": error.to_string() })));
            return;
        }
    };
    let provider_id = header_string(&hello, "providerId");
    let channel_id = header_string(&hello, "channelId");
    let protocol = header_string(&hello, "protocol");
    if provider_id.is_empty() || channel_id.is_empty() || protocol != SSH_DATA_PROTOCOL {
        state.log(
            "warn",
            "ssh_data_hello_rejected",
            Some(json!({ "providerId": provider_id, "channelId": channel_id, "protocol": protocol, "peer": peer_addr.map(|addr| addr.to_string()) })),
        );
        return;
    }
    if !state.active_providers.read().await.contains_key(&provider_id) {
        state.log(
            "warn",
            "ssh_data_provider_not_registered",
            Some(json!({ "providerId": provider_id, "channelId": channel_id })),
        );
        return;
    }
    let (tx, mut rx) = mpsc::unbounded_channel::<SshDataFrame>();
    state.active_ssh_data_channels.lock().await.insert(
        channel_key(&provider_id, &channel_id),
        SshDataChannel {
            provider_id: provider_id.clone(),
            channel_id: channel_id.clone(),
            writer: tx,
            active_session_id: None,
            connected_at_millis: chrono::Utc::now().timestamp_millis(),
        },
    );
    state.log(
        "info",
        "ssh_data_channel_ready",
        Some(json!({ "providerId": provider_id, "channelId": channel_id, "peer": peer_addr.map(|addr| addr.to_string()) })),
    );

    let writer_state = state.clone();
    let writer_provider_id = provider_id.clone();
    let writer_channel_id = channel_id.clone();
    let writer_task = tokio::spawn(async move {
        while let Some(frame) = rx.recv().await {
            if let Err(error) = write_data_frame(&mut writer, frame).await {
                writer_state.log(
                    "warn",
                    "ssh_data_write_failed",
                    Some(json!({ "providerId": writer_provider_id, "channelId": writer_channel_id, "error": error.to_string() })),
                );
                break;
            }
        }
    });

    loop {
        match read_data_frame(&mut reader).await {
            Ok(Some((header, payload))) => {
                forward_provider_frame(&state, &provider_id, &channel_id, header, payload).await;
            }
            Ok(None) => break,
            Err(error) => {
                state.log(
                    "warn",
                    "ssh_data_read_failed",
                    Some(json!({ "providerId": provider_id, "channelId": channel_id, "error": error.to_string() })),
                );
                break;
            }
        }
    }
    state
        .active_ssh_data_channels
        .lock()
        .await
        .remove(&channel_key(&provider_id, &channel_id));
    writer_task.abort();
    state.log(
        "warn",
        "ssh_data_channel_closed",
        Some(json!({ "providerId": provider_id, "channelId": channel_id })),
    );
}

pub async fn spawn_ssh_data_listener(
    state: Arc<AppState>,
    listener: TcpListener,
) -> std::io::Result<()> {
    loop {
        let (stream, _) = listener.accept().await?;
        tokio::spawn(handle_ssh_data_stream(state.clone(), stream));
    }
}

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

    fn test_channel(active_session_id: &str) -> SshDataChannel {
        let (writer, _reader) = mpsc::unbounded_channel();
        SshDataChannel {
            provider_id: "provider".to_string(),
            channel_id: "channel".to_string(),
            writer,
            active_session_id: Some(active_session_id.to_string()),
            connected_at_millis: 1,
        }
    }

    #[test]
    fn stale_session_cannot_release_reassigned_channel() {
        let mut channel = test_channel("new-session");

        assert!(!release_channel_for_session(&mut channel, "old-session"));
        assert_eq!(channel.active_session_id.as_deref(), Some("new-session"));
        assert!(release_channel_for_session(&mut channel, "new-session"));
        assert_eq!(channel.active_session_id, None);
    }
}
