use anyhow::{Context, Result, anyhow};
use clap::{Parser, Subcommand};
use dosh::auth::{
    build_attach_ticket, build_bootstrap, encode_bootstrap, load_or_create_server_secret, now_secs,
    open_attach_ticket, verify_bootstrap,
};
use dosh::config::{ServerConfig, load_server_config};
use dosh::crypto;
use dosh::native::{
    ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello, derive_native_session_key,
    generate_native_ephemeral, host_public_key, host_public_key_line, load_or_create_host_key,
    sign_server_hello, verify_native_user_auth_from_config,
};
use dosh::protocol::{
    self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
    NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody,
    NativeUserAuthBody, PacketKind, ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT,
    StreamClose, StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, 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::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, 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: Option<String>,
        #[arg(long)]
        host_key: bool,
        #[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,
            host_key,
            session,
            mode,
            size,
            client_version: _,
            udp_host,
        } => {
            anyhow::ensure!(protocol == 1, "unsupported protocol {protocol}");
            let config = load_server_config(None)?;
            if host_key {
                let signing_key = load_or_create_host_key(&config)?;
                println!("{}", host_public_key_line(&host_public_key(&signing_key)));
                return Ok(());
            }
            let secret = load_or_create_server_secret(&config)?;
            let nonce = parse_nonce(nonce.as_deref().context("--nonce is required")?)?;
            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 cleanup_state = Arc::clone(&state);
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(Duration::from_secs(5));
        loop {
            interval.tick().await;
            cleanup_disconnected_clients(&cleanup_state);
        }
    });

    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>,
    pending_native: HashMap<[u8; 16], PendingNativeAuth>,
    next_server_stream_id: u64,
}

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>,
    allowed_forwardings: Vec<ForwardingRequest>,
    stream_writers: HashMap<u64, mpsc::Sender<Vec<u8>>>,
}

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

struct PendingNativeAuth {
    client: dosh::native::NativeClientHello,
    server: NativeServerHello,
    session_key: [u8; 32],
    peer: SocketAddr,
    created_at: Instant,
}

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

    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::NativeClientHello => {
            handle_native_client_hello(state, socket, peer, packet.body).await
        }
        PacketKind::NativeUserAuth => handle_native_user_auth(state, socket, peer, &packet).await,
        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
        | PacketKind::Resize
        | PacketKind::Ping
        | PacketKind::Ack
        | PacketKind::StreamOpen
        | PacketKind::StreamOpenOk
        | PacketKind::StreamOpenReject
        | PacketKind::StreamData
        | PacketKind::StreamClose
        | PacketKind::StreamEof
        | PacketKind::StreamWindowAdjust
            if find_client_key(state, &packet.header.conn_id).is_err() =>
        {
            send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client").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,
        PacketKind::StreamOpen => handle_stream_open(state, socket, peer, &packet).await,
        PacketKind::StreamOpenOk => handle_stream_open_ok(state, peer, &packet).await,
        PacketKind::StreamOpenReject => handle_stream_open_reject(state, peer, &packet).await,
        PacketKind::StreamData => handle_stream_data(state, peer, &packet).await,
        PacketKind::StreamClose => handle_stream_close(state, &packet).await,
        _ => Ok(()),
    }
}

async fn handle_native_client_hello(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    body: Vec<u8>,
) -> Result<()> {
    let req: NativeClientHelloBody = protocol::from_body(&body)?;
    let built: Result<([u8; 16], NativeServerHello)> = {
        let mut locked = state.lock().expect("server state poisoned");
        if !locked.config.native_auth {
            Err(anyhow!("native auth disabled"))
        } else if !req
            .hello
            .supported_aead
            .iter()
            .any(|algorithm| algorithm == "chacha20poly1305")
        {
            Err(anyhow!("native auth requires chacha20poly1305"))
        } else if !req
            .hello
            .supported_user_key_algorithms
            .iter()
            .any(|algorithm| algorithm == "ssh-ed25519")
        {
            Err(anyhow!("native auth requires ssh-ed25519"))
        } else {
            let current_user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
            if req.hello.requested_user != current_user {
                Err(anyhow!("native auth user mismatch"))
            } else {
                let host_signing = load_or_create_host_key(&locked.config)?;
                let (server_secret, server_public) = generate_native_ephemeral();
                let mut hello = NativeServerHello {
                    protocol_version: dosh::native::NATIVE_PROTOCOL_VERSION,
                    server_random: crypto::random_32(),
                    server_ephemeral_public: server_public,
                    host_key: host_public_key(&host_signing),
                    chosen_aead: "chacha20poly1305".to_string(),
                    server_key_epoch: 1,
                    auth_challenge: crypto::random_32(),
                    rate_limit_remaining: Some(locked.config.native_auth_rate_limit_per_minute),
                    host_signature: Vec::new(),
                };
                sign_server_hello(&host_signing, &req.hello, &mut hello)?;
                let session_key = derive_native_session_key(
                    &server_secret,
                    req.hello.client_ephemeral_public,
                    &req.hello,
                    &hello,
                )?;
                let mut pending_id = [0u8; 16];
                pending_id.copy_from_slice(&hello.auth_challenge[..16]);
                locked.pending_native.insert(
                    pending_id,
                    PendingNativeAuth {
                        client: req.hello,
                        server: hello.clone(),
                        session_key,
                        peer,
                        created_at: Instant::now(),
                    },
                );
                Ok((pending_id, hello))
            }
        }
    };
    let (pending_id, hello) = match built {
        Ok(built) => built,
        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
    };
    let body = protocol::to_body(&NativeServerHelloBody { hello })?;
    let out = protocol::encode_plain(PacketKind::NativeServerHello, pending_id, 1, 0, &body)?;
    socket.send_to(&out, peer).await?;
    Ok(())
}

async fn handle_native_user_auth(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    peer: SocketAddr,
    packet: &protocol::Packet,
) -> Result<()> {
    let pending = {
        let mut locked = state.lock().expect("server state poisoned");
        locked.pending_native.remove(&packet.header.conn_id)
    };
    let Some(pending) = pending else {
        return send_reject_to_client(socket, peer, packet.header.conn_id, "unknown native auth")
            .await;
    };
    if pending.peer != peer {
        return send_reject_to_client(
            socket,
            peer,
            packet.header.conn_id,
            "native auth peer changed",
        )
        .await;
    }
    if pending.created_at.elapsed() > Duration::from_secs(30) {
        return send_reject_to_client(socket, peer, packet.header.conn_id, "native auth expired")
            .await;
    }
    let body = protocol::decrypt_body(packet, &pending.session_key, CLIENT_TO_SERVER)?;
    let req: NativeUserAuthBody = protocol::from_body(&body)?;
    if pending.client.requested_mode == "doctor" {
        let check = {
            let locked = state.lock().expect("server state poisoned");
            if let Err(err) = verify_native_user_auth_from_config(
                &locked.config,
                &pending.client,
                &pending.server,
                &req.auth,
                Some(peer.ip()),
            )
            .context("verify native user auth")
            {
                return send_reject_to_client(
                    socket,
                    peer,
                    packet.header.conn_id,
                    &format!("{err:#}"),
                )
                .await;
            }
            NativeAuthCheckOkBody {
                requested_user: pending.client.requested_user.clone(),
                native_auth_enabled: locked.config.native_auth,
                allow_tcp_forwarding: locked.config.allow_tcp_forwarding,
                allow_remote_forwarding: locked.config.allow_remote_forwarding,
                allow_agent_forwarding: locked.config.allow_agent_forwarding,
                policy_flags: Vec::new(),
            }
        };
        let body = protocol::to_body(&check)?;
        let out = protocol::encode_encrypted(
            PacketKind::NativeAuthCheckOk,
            packet.header.conn_id,
            1,
            packet.header.seq,
            &pending.session_key,
            SERVER_TO_CLIENT,
            &body,
        )?;
        socket.send_to(&out, peer).await?;
        return Ok(());
    }

    let attached: Result<(
        [u8; 16],
        [u8; 16],
        Vec<u8>,
        [u8; 32],
        String,
        String,
        u64,
        Vec<u8>,
        Vec<ForwardingRequest>,
    )> = {
        let mut locked = state.lock().expect("server state poisoned");
        verify_native_user_auth_from_config(
            &locked.config,
            &pending.client,
            &pending.server,
            &req.auth,
            Some(peer.ip()),
        )
        .context("verify native user auth")
        .and_then(|_| {
            let session_name = pending.client.requested_session.clone();
            let mode = pending.client.requested_mode.clone();
            let cols = pending.client.terminal_size.0;
            let rows = pending.client.terminal_size.1;
            if !locked.sessions.contains_key(&session_name) {
                if locked.config.create_on_attach {
                    locked.ensure_session(&session_name, cols, rows)?;
                } else {
                    return Err(anyhow!("session does not exist"));
                }
            }
            let session_key = pending.session_key;
            let key_id = protocol::session_key_id(&session_key);
            let attach_ticket_psk = crypto::random_32();
            let issued_at = now_secs()?;
            let attach_ticket = build_attach_ticket(
                &locked.secret,
                crypto::sha256(&locked.secret),
                1,
                pending.client.requested_user.clone(),
                session_name.clone(),
                mode.clone(),
                issued_at,
                issued_at + locked.config.attach_ticket_ttl_secs,
                &attach_ticket_psk,
            )?;
            let session = locked
                .sessions
                .get_mut(&session_name)
                .expect("session exists");
            if mode != "view-only" {
                session.pty.resize(cols, rows)?;
                session.parser.set_size(rows, cols);
            }
            let client_id = crypto::random_16();
            let snapshot = screen_snapshot(session.parser.screen());
            let screen = session.parser.screen().clone();
            let output_seq = session.output_seq;
            session.clients.insert(
                client_id,
                ClientState {
                    endpoint: peer,
                    mode: mode.clone(),
                    session_key,
                    last_acked: output_seq,
                    replay: ReplayWindow::default(),
                    send_seq: 1,
                    cols,
                    rows,
                    last_seen: Instant::now(),
                    pending: VecDeque::new(),
                    last_screen: Some(screen),
                    allowed_forwardings: req.auth.requested_forwardings.clone(),
                    stream_writers: HashMap::new(),
                },
            );
            Ok((
                client_id,
                key_id,
                attach_ticket,
                attach_ticket_psk,
                session_name,
                mode,
                output_seq,
                snapshot,
                req.auth.requested_forwardings.clone(),
            ))
        })
    };

    let (
        client_id,
        key_id,
        attach_ticket,
        attach_ticket_psk,
        session_name,
        mode,
        output_seq,
        snapshot,
        requested_forwardings,
    ) = match attached {
        Ok(attached) => attached,
        Err(err) => {
            return send_reject_to_client(socket, peer, packet.header.conn_id, &format!("{err:#}"))
                .await;
        }
    };
    if let Err(err) = start_remote_forwards(
        state,
        socket,
        client_id,
        requested_forwardings.clone(),
        pending.client.requested_session.clone(),
    )
    .await
    {
        remove_client(state, client_id);
        return send_reject_to_client(socket, peer, packet.header.conn_id, &err.to_string()).await;
    }

    let ok = NativeAuthOk {
        client_id,
        session: session_name,
        mode,
        session_key: pending.session_key,
        session_key_id: key_id,
        attach_ticket,
        attach_ticket_psk,
        initial_seq: output_seq,
        snapshot,
        policy_flags: Vec::new(),
    };
    let body = protocol::to_body(&NativeAuthOkBody { ok })?;
    let out = protocol::encode_encrypted(
        PacketKind::NativeAuthOk,
        client_id,
        1,
        packet.header.seq,
        &pending.session_key,
        SERVER_TO_CLIENT,
        &body,
    )?;
    socket.send_to(&out, peer).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");
        if req.bootstrap.mode != "view-only" {
            session.pty.resize(req.cols, req.rows)?;
            session.parser.set_size(req.rows, req.cols);
        }
        let client_id = crypto::random_16();
        let snapshot = screen_snapshot(session.parser.screen());
        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),
                allowed_forwardings: Vec::new(),
                stream_writers: HashMap::new(),
            },
        );
        (
            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 = protocol::session_key_id(&session_key);
    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");
        if req.mode != "view-only" {
            session.pty.resize(req.cols, req.rows)?;
            session.parser.set_size(req.rows, req.cols);
        }
        let client_id = crypto::random_16();
        let snapshot = screen_snapshot(session.parser.screen());
        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),
                allowed_forwardings: Vec::new(),
                stream_writers: HashMap::new(),
            },
        );
        (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<()> {
    send_reject_to_client(socket, peer, [0u8; 16], reason).await
}

async fn send_reject_to_client(
    socket: &UdpSocket,
    peer: SocketAddr,
    client_id: [u8; 16],
    reason: &str,
) -> Result<()> {
    let body = protocol::to_body(&AttachReject {
        reason: reason.to_string(),
    })?;
    let out = protocol::encode_plain(PacketKind::AttachReject, client_id, 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;
        if client.mode != "view-only" {
            session.pty.resize(req.cols, req.rows)?;
            session.parser.set_size(req.rows, req.cols);
        }
        let snapshot = screen_snapshot(session.parser.screen());
        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,
        closed: false,
    };
    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(())
}

fn remove_client(state: &Arc<Mutex<ServerState>>, client_id: [u8; 16]) {
    let mut locked = state.lock().expect("server state poisoned");
    for session in locked.sessions.values_mut() {
        session.clients.remove(&client_id);
    }
}

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 handle_stream_open(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    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 open: StreamOpen = protocol::from_body(&body)?;
    let allowed = {
        let mut locked = state.lock().expect("server state poisoned");
        let config = locked.config.clone();
        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(());
        }
        client.endpoint = peer;
        client.last_seen = Instant::now();
        stream_open_allowed(&config, client, &open)
    };
    if let Err(err) = allowed {
        return send_stream_open_reject(
            state,
            socket,
            packet.header.conn_id,
            open.stream_id,
            err.to_string(),
        )
        .await;
    }

    let target = format!("{}:{}", open.target_host, open.target_port);
    let stream = match TcpStream::connect((open.target_host.as_str(), open.target_port)).await {
        Ok(stream) => stream,
        Err(err) => {
            return send_stream_open_reject(
                state,
                socket,
                packet.header.conn_id,
                open.stream_id,
                format!("connect {target}: {err}"),
            )
            .await;
        }
    };
    let (mut reader, mut writer) = stream.into_split();
    let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
    {
        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"))?;
        client.stream_writers.insert(open.stream_id, writer_tx);
    }

    tokio::spawn(async move {
        while let Some(bytes) = writer_rx.recv().await {
            if writer.write_all(&bytes).await.is_err() {
                break;
            }
        }
    });

    send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;

    let state = Arc::clone(state);
    let socket = Arc::clone(socket);
    let client_id = packet.header.conn_id;
    let stream_id = open.stream_id;
    tokio::spawn(async move {
        let mut buf = [0u8; 16 * 1024];
        loop {
            match reader.read(&mut buf).await {
                Ok(0) => break,
                Ok(n) => {
                    if let Err(err) = send_stream_data_to_client(
                        &state,
                        &socket,
                        client_id,
                        stream_id,
                        buf[..n].to_vec(),
                    )
                    .await
                    {
                        eprintln!("stream send error: {err:#}");
                        break;
                    }
                }
                Err(_) => break,
            }
        }
        let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
    });

    Ok(())
}

async fn handle_stream_data(
    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 data: StreamData = protocol::from_body(&body)?;
    let writer = {
        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(());
        }
        client.endpoint = peer;
        client.last_seen = Instant::now();
        client.stream_writers.get(&data.stream_id).cloned()
    };
    if let Some(writer) = writer {
        let _ = writer.send(data.bytes).await;
    }
    Ok(())
}

async fn handle_stream_open_ok(
    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 _ok: StreamOpenOk = 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(());
    }
    client.endpoint = peer;
    client.last_seen = Instant::now();
    Ok(())
}

async fn handle_stream_open_reject(
    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 reject: StreamOpenReject = 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(());
    }
    client.endpoint = peer;
    client.last_seen = Instant::now();
    client.stream_writers.remove(&reject.stream_id);
    Ok(())
}

async fn handle_stream_close(
    state: &Arc<Mutex<ServerState>>,
    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 close: StreamClose = protocol::from_body(&body)?;
    let mut locked = state.lock().expect("server state poisoned");
    let Some(session) = locked.sessions.get_mut(&session_name) else {
        return Ok(());
    };
    let Some(client) = session.clients.get_mut(&packet.header.conn_id) else {
        return Ok(());
    };
    if !client.replay.accept(packet.header.seq) {
        return Ok(());
    }
    client.last_seen = Instant::now();
    client.stream_writers.remove(&close.stream_id);
    Ok(())
}

async fn start_remote_forwards(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    requested_forwardings: Vec<ForwardingRequest>,
    session_name: String,
) -> Result<()> {
    let remote_forwards = requested_forwardings
        .into_iter()
        .filter(|forward| forward.kind == ForwardingKind::Remote)
        .collect::<Vec<_>>();
    if remote_forwards.is_empty() {
        return Ok(());
    }
    let config = {
        let locked = state.lock().expect("server state poisoned");
        locked.config.clone()
    };
    anyhow::ensure!(
        config.allow_remote_forwarding,
        "remote forwarding disabled on server"
    );
    for forward in remote_forwards {
        let bind_host = forward
            .bind_host
            .clone()
            .unwrap_or_else(|| "127.0.0.1".to_string());
        anyhow::ensure!(
            is_loopback_bind(&bind_host),
            "remote forwarding may only bind loopback in this server build"
        );
        let listen_port = forward.listen_port;
        let target_host = forward
            .target_host
            .clone()
            .ok_or_else(|| anyhow!("remote forward target host is required"))?;
        let target_port = forward
            .target_port
            .ok_or_else(|| anyhow!("remote forward target port is required"))?;
        let bind = format!("{bind_host}:{listen_port}");
        let listener = TcpListener::bind(&bind)
            .await
            .with_context(|| format!("bind remote forward {bind}"))?;
        let state = Arc::clone(state);
        let socket = Arc::clone(socket);
        let session_name = session_name.clone();
        tokio::spawn(async move {
            loop {
                let Ok((stream, _)) = listener.accept().await else {
                    break;
                };
                let stream_id = allocate_server_stream_id(&state);
                let (mut reader, mut writer) = stream.into_split();
                let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
                {
                    let mut locked = state.lock().expect("server state poisoned");
                    if let Some(session) = locked.sessions.get_mut(&session_name)
                        && let Some(client) = session.clients.get_mut(&client_id)
                    {
                        client.stream_writers.insert(stream_id, writer_tx);
                    } else {
                        break;
                    }
                }
                tokio::spawn(async move {
                    while let Some(bytes) = writer_rx.recv().await {
                        if writer.write_all(&bytes).await.is_err() {
                            break;
                        }
                    }
                });
                if let Err(err) = send_stream_open_to_client(
                    &state,
                    &socket,
                    client_id,
                    stream_id,
                    target_host.clone(),
                    target_port,
                )
                .await
                {
                    eprintln!("remote forward open error: {err:#}");
                    break;
                }
                let state = Arc::clone(&state);
                let socket = Arc::clone(&socket);
                tokio::spawn(async move {
                    let mut buf = [0u8; 16 * 1024];
                    loop {
                        match reader.read(&mut buf).await {
                            Ok(0) => break,
                            Ok(n) => {
                                if let Err(err) = send_stream_data_to_client(
                                    &state,
                                    &socket,
                                    client_id,
                                    stream_id,
                                    buf[..n].to_vec(),
                                )
                                .await
                                {
                                    eprintln!("remote forward data error: {err:#}");
                                    break;
                                }
                            }
                            Err(_) => break,
                        }
                    }
                    let _ =
                        send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
                });
            }
        });
    }
    Ok(())
}

fn allocate_server_stream_id(state: &Arc<Mutex<ServerState>>) -> u64 {
    let mut locked = state.lock().expect("server state poisoned");
    let stream_id = locked.next_server_stream_id;
    locked.next_server_stream_id = locked.next_server_stream_id.wrapping_add(1).max(1u64 << 63);
    stream_id
}

fn is_loopback_bind(host: &str) -> bool {
    matches!(host, "127.0.0.1" | "localhost" | "::1")
}

fn stream_open_allowed(
    config: &ServerConfig,
    client: &ClientState,
    open: &StreamOpen,
) -> Result<()> {
    anyhow::ensure!(config.allow_tcp_forwarding, "TCP forwarding disabled");
    let allowed = client
        .allowed_forwardings
        .iter()
        .any(|forward| match forward.kind {
            ForwardingKind::Local => {
                forward.target_host.as_deref() == Some(open.target_host.as_str())
                    && forward.target_port == Some(open.target_port)
            }
            ForwardingKind::Dynamic => true,
            ForwardingKind::Remote => false,
        });
    anyhow::ensure!(
        allowed,
        "stream target {}:{} was not requested during auth",
        open.target_host,
        open.target_port
    );
    Ok(())
}

async fn send_stream_open_ok(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
) -> Result<()> {
    let body = protocol::to_body(&StreamOpenOk { stream_id })?;
    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenOk, body).await
}

async fn send_stream_open_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    target_host: String,
    target_port: u16,
) -> Result<()> {
    let body = protocol::to_body(&StreamOpen {
        stream_id,
        target_host,
        target_port,
    })?;
    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpen, body).await
}

async fn send_stream_open_reject(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    reason: String,
) -> Result<()> {
    let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?;
    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenReject, body).await
}

async fn send_stream_data_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
    bytes: Vec<u8>,
) -> Result<()> {
    let body = protocol::to_body(&StreamData { stream_id, bytes })?;
    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamData, body).await
}

async fn send_stream_close_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    stream_id: u64,
) -> Result<()> {
    {
        let mut locked = state.lock().expect("server state poisoned");
        for session in locked.sessions.values_mut() {
            if let Some(client) = session.clients.get_mut(&client_id) {
                client.stream_writers.remove(&stream_id);
                break;
            }
        }
    }
    let body = protocol::to_body(&StreamClose { stream_id })?;
    send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamClose, body).await
}

async fn send_stream_packet_to_client(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    client_id: [u8; 16],
    kind: PacketKind,
    body: Vec<u8>,
) -> Result<()> {
    let send = {
        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(&client_id) {
                client.send_seq += 1;
                let packet = protocol::encode_encrypted(
                    kind,
                    client_id,
                    client.send_seq,
                    client.last_acked,
                    &client.session_key,
                    SERVER_TO_CLIENT,
                    &body,
                )?;
                found = Some((client.endpoint, packet));
                break;
            }
        }
        found
    };
    if let Some((endpoint, packet)) = send {
        socket.send_to(&packet, endpoint).await?;
    }
    Ok(())
}

async fn broadcast_output(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    output: PtyOutput,
) -> Result<()> {
    if output.exited {
        return broadcast_session_exit(state, socket, output.session).await;
    }

    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 (bytes, snapshot) = if client.pending.len() >= retransmit_window {
                client.pending.clear();
                (screen_snapshot(&current_screen), true)
            } else {
                (output.bytes.clone(), false)
            };
            let frame = Frame {
                session: output.session.clone(),
                output_seq,
                bytes,
                snapshot,
                closed: false,
            };
            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 broadcast_session_exit(
    state: &Arc<Mutex<ServerState>>,
    socket: &Arc<UdpSocket>,
    session_name: String,
) -> Result<()> {
    let sends = {
        let mut locked = state.lock().expect("server state poisoned");
        let Some(mut session) = locked.sessions.remove(&session_name) else {
            return Ok(());
        };
        session.output_seq += 1;
        let output_seq = session.output_seq;
        let mut sends = Vec::new();
        for (client_id, mut client) in session.clients.drain() {
            client.send_seq += 1;
            let frame = Frame {
                session: session_name.clone(),
                output_seq,
                bytes: b"\r\n[dosh session exited]\r\n".to_vec(),
                snapshot: false,
                closed: true,
            };
            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,
            )?;
            sends.push((client.endpoint, packet));
        }
        sends
    };
    for (endpoint, packet) in sends {
        socket.send_to(&packet, endpoint).await?;
    }
    Ok(())
}

fn screen_snapshot(screen: &vt100::Screen) -> Vec<u8> {
    let mut bytes = Vec::new();
    if screen.alternate_screen() {
        bytes.extend_from_slice(b"\x1b[?1049h");
    } else {
        bytes.extend_from_slice(b"\x1b[?1049l");
    }
    bytes.extend(screen.state_formatted());
    bytes
}

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 cleanup_disconnected_clients(state: &Arc<Mutex<ServerState>>) {
    let mut locked = state.lock().expect("server state poisoned");
    let timeout = Duration::from_secs(locked.config.client_timeout_secs.max(1));
    let now = Instant::now();
    for session in locked.sessions.values_mut() {
        session
            .clients
            .retain(|_, client| now.duration_since(client.last_seen) <= timeout);
    }
}

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()?))
}