//! VNC screen websocket relay + agent icon.

use axum::{
    extract::State,
    http::StatusCode,
    response::{IntoResponse, Response},
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

use super::AppState;

/// This agent's icon. Serves the operator-configured SVG from
/// `/etc/hyperhive/icon.svg` (set via the `hyperhive.icon` agent.nix
/// option) when present, otherwise the bundled default hyperhive logo.
/// Always returns an image, so consumers (dashboard, favicon) can hit
/// `/icon` unconditionally without probing whether one is configured.
pub(super) async fn serve_icon() -> impl IntoResponse {
    // Per-agent icon overrides go through `/etc/hyperhive/icon.svg`
    // (set via the `hyperhive.icon` agent.nix option); the bundled
    // default is resolved at runtime from
    // `$HIVE_ASSETS_DIR/branding/hyperhive.svg`. If neither file can
    // be read we serve an empty body — keeps the response a valid SVG
    // content-type without a panic on a misconfigured container.
    let body = std::fs::read_to_string("/etc/hyperhive/icon.svg").unwrap_or_else(|_| {
        std::fs::read_to_string(hive_sh4re::assets::branding_svg()).unwrap_or_default()
    });
    ([("content-type", "image/svg+xml")], body)
}

/// WebSocket handler: upgrade then pump bytes between the WS client and
/// the VNC server on `127.0.0.1:<vnc_port>`. Returns 404 when gui is not
/// enabled for this agent.
pub(super) async fn screen_ws(
    ws: axum::extract::ws::WebSocketUpgrade,
    State(state): State<AppState>,
) -> Response {
    let Some(vnc_port) = state.gui_vnc_port else {
        return (StatusCode::NOT_FOUND, "gui not enabled for this agent").into_response();
    };
    ws.on_upgrade(move |socket| relay_ws_vnc(socket, vnc_port))
}

/// Pure byte pump: forwards raw bytes between the WebSocket client and
/// the VNC TCP stream. Transparent to any RFB variant (plain, `VeNCrypt`).
async fn relay_ws_vnc(socket: axum::extract::ws::WebSocket, vnc_port: u16) {
    // Import futures traits locally so they don't conflict with
    // tokio_stream::StreamExt used at module scope.
    use axum::extract::ws::Message;
    use futures_util::{SinkExt, StreamExt as _};

    let addr = format!("127.0.0.1:{vnc_port}");
    let Ok(tcp) = tokio::net::TcpStream::connect(&addr).await else {
        tracing::warn!(%addr, "screen/ws: could not connect to VNC server");
        return;
    };
    let (mut tcp_rx, mut tcp_tx) = tcp.into_split();
    let (mut ws_tx, mut ws_rx) = socket.split();

    // WS → TCP
    let ws_to_tcp = tokio::spawn(async move {
        while let Some(Ok(msg)) = futures_util::StreamExt::next(&mut ws_rx).await {
            match msg {
                Message::Binary(data) if tcp_tx.write_all(&data).await.is_err() => {
                    break;
                }
                Message::Close(_) => break,
                _ => {} // ping/pong/text: ignore
            }
        }
    });

    // TCP → WS
    let tcp_to_ws = tokio::spawn(async move {
        let mut buf = vec![0u8; 8192];
        loop {
            match tcp_rx.read(&mut buf).await {
                Ok(0) | Err(_) => break,
                Ok(n) => {
                    if ws_tx
                        .send(Message::Binary(buf[..n].to_vec().into()))
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
            }
        }
    });

    // Wait for either direction to close, then let both tasks drop.
    tokio::select! {
        _ = ws_to_tcp => {}
        _ = tcp_to_ws => {}
    }
}