use anyhow::{Context, Result, anyhow};
use clap::{Parser, Subcommand};
use dosh::auth::{
    build_bootstrap, encode_bootstrap, load_or_create_server_secret, open_attach_ticket,
    verify_bootstrap,
};
use dosh::config::{ServerConfig, load_server_config};
use dosh::crypto;
use dosh::protocol::{
    self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
    PacketKind, ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody,
    TicketAttachEnvelope, TicketAttachOkEnvelope,
};
use dosh::pty::{PtyHandle, PtyOutput, spawn_pty_session};
use std::collections::{HashMap, VecDeque};
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::net::UdpSocket;
use tokio::sync::mpsc;

#[derive(Debug, Parser)]
#[command(name = "dosh-server")]
struct Args {
    #[command(subcommand)]
    command: Command,
}

#[derive(Debug, Subcommand)]
enum Command {
    Serve {
        #[arg(long)]
        config: Option<std::path::PathBuf>,
    },
    Auth {
        #[arg(long, default_value_t = 1)]
        protocol: u8,
        #[arg(long)]
        nonce: String,
        #[arg(long, default_value = "default")]
        session: String,
        #[arg(long, default_value = "read-write")]
        mode: String,
        #[arg(long, default_value = "80x24")]
        size: String,
        #[arg(long, default_value = "dev")]
        client_version: String,
        #[arg(long)]
        udp_host: Option<String>,
    },
}

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    let args = Args::parse();
    match args.command {
        Command::Serve { config } => serve(config).await,
        Command::Auth {
            protocol,
            nonce,
            session,
            mode,
            size,
            client_version: _,
            udp_host,
        } => {
            anyhow::ensure!(protocol == 1, "unsupported protocol {protocol}");
            let config = load_server_config(None)?;
            let secret = load_or_create_server_secret(&config)?;
            let nonce = parse_nonce(&nonce)?;
            let size = parse_size(&size)?;
            let user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
            let udp_host = udp_host.unwrap_or_else(|| "127.0.0.1".to_string());
            let resp =
                build_bootstrap(&config, &secret, user, session, mode, size, nonce, udp_host)?;
            println!("{}", encode_bootstrap(&resp)?);
            Ok(())
        }
    }
}

async fn serve(config_path: Option<std::path::PathBuf>) -> Result<()> {
    let config = load_server_config(config_path)?;
    let secret = load_or_create_server_secret(&config)?;
    let bind = format!("{}:{}", config.bind, config.port);
    let socket = Arc::new(
        UdpSocket::bind(&bind)
            .await
            .with_context(|| format!("bind {bind}"))?,
    );
    eprintln!("dosh-server listening on {bind}");

    let (pty_tx, mut pty_rx) = mpsc::unbounded_channel();
    let state = Arc::new(Mutex::new(ServerState::new(
        config.clone(),
        secret,
        pty_tx.clone(),
    )));
    {
        let mut locked = state.lock().expect("server state poisoned");
        for session in config.prewarm_sessions.clone() {
            locked.ensure_session(&session, 80, 24)?;
        }
    }

    let output_state = Arc::clone(&state);
    let output_socket = Arc::clone(&socket);
    tokio::spawn(async move {
        while let Some(output) = pty_rx.recv().await {
            if let Err(err) = broadcast_output(&output_state, &output_socket, output).await {
                eprintln!("broadcast error: {err:#}");
            }
        }
    });

    let retransmit_state = Arc::clone(&state);
    let retransmit_socket = Arc::clone(&socket);
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(Duration::from_millis(100));
        loop {
            interval.tick().await;
            if let Err(err) = retransmit_pending(&retransmit_state, &retransmit_socket).await {
                eprintln!("retransmit error: {err:#}");
            }
        }
    });

    let mut buf = vec![0u8; 65535];
    loop {
        let (n, peer) = socket.recv_from(&mut buf).await?;
        if let Err(err) = handle_packet(&state, &socket, peer, &buf[..n]).await {
            eprintln!("packet from {peer}: {err:#}");
        }
    }
}

struct ServerState {
    config: ServerConfig,
    secret: [u8; 32],
    pty_tx: mpsc::UnboundedSender<PtyOutput>,
    sessions: HashMap<String, Session>,
}

struct Session {
    pty: PtyHandle,
    parser: vt100::Parser,
    clients: HashMap<[u8; 16], ClientState>,
    output_seq: u64,
    recent: VecDeque<Vec<u8>>,
}

#[derive(Clone)]
struct ClientState {
    endpoint: SocketAddr,
    mode: String,
    session_key: [u8; 32],
    last_acked: u64,
    replay: ReplayWindow,
    send_seq: u64,
    cols: u16,
    rows: u16,
    last_seen: Instant,
    pending: VecDeque<PendingFrame>,
    last_screen: Option<vt100::Screen>,
}

#[derive(Clone)]
struct PendingFrame {
    output_seq: u64,
    packet: Vec<u8>,
    last_sent: Instant,
    attempts: u8,
}

impl ServerState {
    fn new(
        config: ServerConfig,
        secret: [u8; 32],
        pty_tx: mpsc::UnboundedSender<PtyOutput>,
    ) -> Self {
        Self {
            config,
            secret,
            pty_tx,
            sessions: HashMap::new(),
        }
    }

    fn ensure_session(&mut self, name: &str, cols: u16, rows: u16) -> Result<()> {
        if self.sessions.contains_key(name) {
            return Ok(());
        }
        let pty = spawn_pty_session(
            name.to_string(),
            &self.config.shell,
            cols.max(1),
            rows.max(1),
            self.pty_tx.clone(),
        )?;
        self.sessions.insert(
            name.to_string(),
            Session {
                pty,
                parser: vt100::Parser::new(rows.max(1), cols.max(1), self.config.scrollback),
                clients: HashMap::new(),
                output_seq: 0,
                recent: VecDeque::with_capacity(self.config.scrollback),
            },
        );
        Ok(())
    }
}

async fn handle_packet(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    raw: &[u8],
) -> Result<()> {
    let packet = protocol::decode(raw)?;
    match packet.header.kind {
        PacketKind::BootstrapAttachRequest => {
            handle_bootstrap_attach(state, socket, peer, packet.body).await
        }
        PacketKind::TicketAttachRequest => {
            handle_ticket_attach(state, socket, peer, packet.body).await
        }
        PacketKind::ResumeRequest => handle_resume(state, socket, peer, &packet).await,
        PacketKind::Input => handle_input(state, peer, &packet).await,
        PacketKind::Resize => handle_resize(state, peer, &packet).await,
        PacketKind::Ping => handle_ping(state, socket, peer, &packet).await,
        PacketKind::Ack => handle_ack(state, &packet).await,
        PacketKind::Detach => handle_detach(state, &packet).await,
        _ => Ok(()),
    }
}

async fn handle_bootstrap_attach(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    body: Vec<u8>,
) -> Result<()> {
    let req: BootstrapAttachRequest = protocol::from_body(&body)?;
    let (client_id, key, key_id, session_name, mode, output_seq, snapshot) = {
        let mut locked = state.lock().expect("server state poisoned");
        if !verify_bootstrap(&req.bootstrap, &locked.secret)? {
            return send_reject(socket, peer, "invalid or expired bootstrap").await;
        }
        if !locked.sessions.contains_key(&req.bootstrap.session) {
            if locked.config.create_on_attach {
                locked.ensure_session(&req.bootstrap.session, req.cols, req.rows)?;
            } else {
                return send_reject(socket, peer, "session does not exist").await;
            }
        }
        let session = locked
            .sessions
            .get_mut(&req.bootstrap.session)
            .expect("session exists");
        let client_id = crypto::random_16();
        let snapshot = session.parser.screen().state_formatted();
        let screen = session.parser.screen().clone();
        let output_seq = session.output_seq;
        session.clients.insert(
            client_id,
            ClientState {
                endpoint: peer,
                mode: req.bootstrap.mode.clone(),
                session_key: req.bootstrap.session_key,
                last_acked: output_seq,
                replay: ReplayWindow::default(),
                send_seq: 1,
                cols: req.cols,
                rows: req.rows,
                last_seen: Instant::now(),
                pending: VecDeque::new(),
                last_screen: Some(screen),
            },
        );
        (
            client_id,
            req.bootstrap.session_key,
            req.bootstrap.session_key_id,
            req.bootstrap.session.clone(),
            req.bootstrap.mode.clone(),
            output_seq,
            snapshot,
        )
    };
    let ok = AttachOk {
        client_id,
        session: session_name,
        mode,
        session_key: key,
        session_key_id: key_id,
        initial_seq: output_seq,
        snapshot,
    };
    let body = protocol::to_body(&ok)?;
    let out = protocol::encode_encrypted(
        PacketKind::AttachOk,
        client_id,
        1,
        0,
        &key,
        SERVER_TO_CLIENT,
        &body,
    )?;
    socket.send_to(&out, peer).await?;
    Ok(())
}

async fn handle_ticket_attach(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    body: Vec<u8>,
) -> Result<()> {
    let env: TicketAttachEnvelope = protocol::from_body(&body)?;
    let (ticket, request_plain) = {
        let locked = state.lock().expect("server state poisoned");
        if !locked.config.allow_attach_tickets {
            return send_reject(socket, peer, "attach tickets disabled").await;
        }
        let ticket = open_attach_ticket(&locked.secret, &env.ticket)?;
        let request_key = crypto::hkdf32(
            &ticket.psk,
            &env.client_nonce,
            b"dosh/ticket-attach-request/v1",
        )?;
        let request_plain = crypto::open(
            &request_key,
            &env.client_nonce,
            b"dosh-ticket-attach-request-v1",
            &env.ciphertext,
        )?;
        (ticket, request_plain)
    };
    let req: TicketAttachBody = protocol::from_body(&request_plain)?;
    if req.session != ticket.session || req.mode != ticket.mode {
        return send_reject(socket, peer, "ticket scope mismatch").await;
    }

    let session_key = crypto::random_32();
    let session_key_id = {
        let digest = crypto::sha256(&session_key);
        let mut out = [0u8; 16];
        out.copy_from_slice(&digest[..16]);
        out
    };
    let (client_id, output_seq, snapshot) = {
        let mut locked = state.lock().expect("server state poisoned");
        if !locked.sessions.contains_key(&req.session) {
            if locked.config.create_on_attach {
                locked.ensure_session(&req.session, req.cols, req.rows)?;
            } else {
                return send_reject(socket, peer, "session does not exist").await;
            }
        }
        let session = locked
            .sessions
            .get_mut(&req.session)
            .expect("session exists");
        let client_id = crypto::random_16();
        let snapshot = session.parser.screen().state_formatted();
        let screen = session.parser.screen().clone();
        let output_seq = session.output_seq;
        session.clients.insert(
            client_id,
            ClientState {
                endpoint: peer,
                mode: req.mode.clone(),
                session_key,
                last_acked: output_seq,
                replay: ReplayWindow::default(),
                send_seq: 1,
                cols: req.cols,
                rows: req.rows,
                last_seen: Instant::now(),
                pending: VecDeque::new(),
                last_screen: Some(screen),
            },
        );
        (client_id, output_seq, snapshot)
    };

    let ok = AttachOk {
        client_id,
        session: req.session,
        mode: req.mode,
        session_key,
        session_key_id,
        initial_seq: output_seq,
        snapshot,
    };
    let ok_plain = protocol::to_body(&ok)?;
    let server_nonce = crypto::random_12();
    let mut salt = Vec::with_capacity(24);
    salt.extend_from_slice(&env.client_nonce);
    salt.extend_from_slice(&server_nonce);
    let response_key = crypto::hkdf32(&ticket.psk, &salt, b"dosh/ticket-attach-ok/v1")?;
    let ciphertext = crypto::seal(
        &response_key,
        &server_nonce,
        b"dosh-ticket-attach-ok-v1",
        &ok_plain,
    )?;
    let envelope = TicketAttachOkEnvelope {
        server_nonce,
        ciphertext,
    };
    let body = protocol::to_body(&envelope)?;
    let out = protocol::encode_plain(PacketKind::AttachOk, client_id, 1, 0, &body)?;
    socket.send_to(&out, peer).await?;
    Ok(())
}

async fn send_reject(socket: &UdpSocket, peer: SocketAddr, reason: &str) -> Result<()> {
    let body = protocol::to_body(&AttachReject {
        reason: reason.to_string(),
    })?;
    let out = protocol::encode_plain(PacketKind::AttachReject, [0u8; 16], 0, 0, &body)?;
    socket.send_to(&out, peer).await?;
    Ok(())
}

async fn handle_resume(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = match find_client_key(state, &packet.header.conn_id) {
        Ok(found) => found,
        Err(_) => return send_reject(socket, peer, "unknown client").await,
    };
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let req: ResumeRequest = protocol::from_body(&body)?;
    let (send_seq, output_seq, snapshot) = {
        let mut locked = state.lock().expect("server state poisoned");
        let session = locked
            .sessions
            .get_mut(&req.session)
            .ok_or_else(|| anyhow!("unknown session"))?;
        let client = session
            .clients
            .get_mut(&packet.header.conn_id)
            .ok_or_else(|| anyhow!("unknown client"))?;
        if !client.replay.accept(packet.header.seq) {
            return Ok(());
        }
        client.endpoint = peer;
        client.last_acked = req.last_rendered_seq;
        client.cols = req.cols;
        client.rows = req.rows;
        client.last_seen = Instant::now();
        client.send_seq += 1;
        let snapshot = session.parser.screen().state_formatted();
        client.last_screen = Some(session.parser.screen().clone());
        (client.send_seq, session.output_seq, snapshot)
    };
    let frame = Frame {
        session: session_name,
        output_seq,
        bytes: snapshot,
        snapshot: true,
    };
    let body = protocol::to_body(&frame)?;
    let out = protocol::encode_encrypted(
        PacketKind::ResumeOk,
        packet.header.conn_id,
        send_seq,
        packet.header.seq,
        &key,
        SERVER_TO_CLIENT,
        &body,
    )?;
    socket.send_to(&out, peer).await?;
    Ok(())
}

async fn handle_input(
    state: &Arc<Mutex<ServerState>>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_key(state, &packet.header.conn_id)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let input: Input = protocol::from_body(&body)?;
    let mut locked = state.lock().expect("server state poisoned");
    let session = locked
        .sessions
        .get_mut(&session_name)
        .ok_or_else(|| anyhow!("unknown session"))?;
    let client = session
        .clients
        .get_mut(&packet.header.conn_id)
        .ok_or_else(|| anyhow!("unknown client"))?;
    if !client.replay.accept(packet.header.seq) {
        return Ok(());
    }
    if client.endpoint != peer {
        client.endpoint = peer;
    }
    client.last_seen = Instant::now();
    if client.mode == "view-only" {
        return Ok(());
    }
    session.pty.write_all(&input.bytes)?;
    Ok(())
}

async fn handle_resize(
    state: &Arc<Mutex<ServerState>>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, session_name) = find_client_key(state, &packet.header.conn_id)?;
    let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let resize: Resize = protocol::from_body(&body)?;
    let mut locked = state.lock().expect("server state poisoned");
    let session = locked
        .sessions
        .get_mut(&session_name)
        .ok_or_else(|| anyhow!("unknown session"))?;
    let client = session
        .clients
        .get_mut(&packet.header.conn_id)
        .ok_or_else(|| anyhow!("unknown client"))?;
    if !client.replay.accept(packet.header.seq) {
        return Ok(());
    }
    if client.mode != "view-only" {
        client.endpoint = peer;
        client.cols = resize.cols;
        client.rows = resize.rows;
        session.pty.resize(resize.cols, resize.rows)?;
        session.parser.set_size(resize.rows, resize.cols);
    }
    Ok(())
}

async fn handle_ping(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let (key, _) = find_client_key(state, &packet.header.conn_id)?;
    let seq = {
        let mut locked = state.lock().expect("server state poisoned");
        let mut found = None;
        for session in locked.sessions.values_mut() {
            if let Some(client) = session.clients.get_mut(&packet.header.conn_id) {
                if !client.replay.accept(packet.header.seq) {
                    return Ok(());
                }
                client.last_seen = Instant::now();
                client.send_seq += 1;
                found = Some(client.send_seq);
                break;
            }
        }
        found.ok_or_else(|| anyhow!("unknown client"))?
    };
    let out = protocol::encode_encrypted(
        PacketKind::Pong,
        packet.header.conn_id,
        seq,
        packet.header.seq,
        &key,
        SERVER_TO_CLIENT,
        b"",
    )?;
    socket.send_to(&out, peer).await?;
    Ok(())
}

async fn handle_detach(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
    let mut locked = state.lock().expect("server state poisoned");
    for session in locked.sessions.values_mut() {
        session.clients.remove(&packet.header.conn_id);
    }
    Ok(())
}

async fn handle_ack(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
    let (key, _) = find_client_key(state, &packet.header.conn_id)?;
    let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
    let mut locked = state.lock().expect("server state poisoned");
    for session in locked.sessions.values_mut() {
        if let Some(client) = session.clients.get_mut(&packet.header.conn_id) {
            if !client.replay.accept(packet.header.seq) {
                return Ok(());
            }
            client.last_seen = Instant::now();
            client.last_acked = packet.header.ack;
            while client
                .pending
                .front()
                .is_some_and(|pending| pending.output_seq <= packet.header.ack)
            {
                client.pending.pop_front();
            }
            return Ok(());
        }
    }
    Ok(())
}

async fn broadcast_output(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    output: PtyOutput,
) -> Result<()> {
    let sends = {
        let mut locked = state.lock().expect("server state poisoned");
        let scrollback = locked.config.scrollback;
        let retransmit_window = locked.config.retransmit_window;
        let session = locked
            .sessions
            .get_mut(&output.session)
            .ok_or_else(|| anyhow!("unknown session"))?;
        session.parser.process(&output.bytes);
        session.output_seq += 1;
        let output_seq = session.output_seq;
        session.recent.push_back(output.bytes.clone());
        while session.recent.len() > scrollback {
            session.recent.pop_front();
        }
        let mut sends = Vec::new();
        for (client_id, client) in session.clients.iter_mut() {
            client.send_seq += 1;
            let current_screen = session.parser.screen().clone();
            let mut snapshot = false;
            let mut bytes = if client.pending.len() >= retransmit_window {
                client.pending.clear();
                snapshot = true;
                current_screen.state_formatted()
            } else if let Some(prev) = &client.last_screen {
                current_screen.state_diff(prev)
            } else {
                snapshot = true;
                current_screen.state_formatted()
            };
            if bytes.is_empty() {
                bytes = output.bytes.clone();
            }
            let frame = Frame {
                session: output.session.clone(),
                output_seq,
                bytes,
                snapshot,
            };
            let body = protocol::to_body(&frame)?;
            let packet = protocol::encode_encrypted(
                PacketKind::Frame,
                *client_id,
                client.send_seq,
                client.last_acked,
                &client.session_key,
                SERVER_TO_CLIENT,
                &body,
            )?;
            while client.pending.len() >= retransmit_window {
                client.pending.pop_front();
            }
            client.last_screen = Some(current_screen);
            client.pending.push_back(PendingFrame {
                output_seq,
                packet: packet.clone(),
                last_sent: Instant::now(),
                attempts: 0,
            });
            sends.push((client.endpoint, packet));
        }
        sends
    };
    for (endpoint, packet) in sends {
        socket.send_to(&packet, endpoint).await?;
    }
    Ok(())
}

async fn retransmit_pending(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
) -> Result<()> {
    let sends = {
        let mut locked = state.lock().expect("server state poisoned");
        let now = Instant::now();
        let mut sends = Vec::new();
        for session in locked.sessions.values_mut() {
            for client in session.clients.values_mut() {
                for pending in client.pending.iter_mut() {
                    if pending.output_seq <= client.last_acked {
                        continue;
                    }
                    if now.duration_since(pending.last_sent) >= Duration::from_millis(200)
                        && pending.attempts < 8
                    {
                        pending.last_sent = now;
                        pending.attempts += 1;
                        sends.push((client.endpoint, pending.packet.clone()));
                    }
                }
            }
        }
        sends
    };
    for (endpoint, packet) in sends {
        socket.send_to(&packet, endpoint).await?;
    }
    Ok(())
}

fn find_client_key(
    state: &Arc<Mutex<ServerState>>,
    client_id: &[u8; 16],
) -> Result<([u8; 32], String)> {
    let locked = state.lock().expect("server state poisoned");
    for (name, session) in &locked.sessions {
        if let Some(client) = session.clients.get(client_id) {
            return Ok((client.session_key, name.clone()));
        }
    }
    Err(anyhow!("unknown client"))
}

fn parse_nonce(raw: &str) -> Result<[u8; 12]> {
    let bytes = base64::Engine::decode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, raw)
        .context("decode nonce")?;
    anyhow::ensure!(bytes.len() == 12, "nonce must decode to 12 bytes");
    let mut out = [0u8; 12];
    out.copy_from_slice(&bytes);
    Ok(out)
}

fn parse_size(raw: &str) -> Result<(u16, u16)> {
    let (cols, rows) = raw.split_once('x').context("size must be COLSxROWS")?;
    Ok((cols.parse()?, rows.parse()?))
}