arcly-http 0.1.1

//! WebSocket boundary: upgrade, per-socket read/write pumps, event dispatch.
//!
//! This is the *only* module that touches `axum::extract::ws` — the analogue of
//! [`crate::web::boundary`] for the real-time layer. Everything above it speaks
//! arcly types ([`WsClient`], [`WsMessage`], [`GatewayRuntime`]).
//!
//! ## Per-connection model (no hot-path locks)
//!
//! ```text
//!            ┌─────────────────── handle_socket task ───────────────────┐
//!  socket ──>│ reader: stream.next() ─> dispatch(event) ─> handler fut  │
//!            │ writer: rx.recv()     ─> sink.send(frame)                │
//!            └───────────────────────────────────────────────────────────┘
//! ```
//! The reader and writer run as independent halves of the split socket. Inbound
//! frames are parsed and routed through the gateway's `dispatch` table (an
//! immutable `&HashMap` — lock-free read). Outbound frames are produced by any
//! task via the registry's sharded channels and drained by this socket's writer.

use std::sync::Arc;

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::http::HeaderMap;
use axum::routing::{get, MethodRouter};
use futures::{SinkExt, StreamExt};
use tokio::sync::{mpsc, oneshot};

use crate::core::engine::FrozenDiContainer;
use crate::realtime::connection::{ConnectionRegistry, WsClient, WsMessage};
use crate::realtime::gateway::GatewayRuntime;
use crate::web::context::Claims;

/// Build the axum `MethodRouter` that upgrades HTTP→WebSocket for one gateway.
///
/// If a `JwtService` has been provided in the DI container, the
/// `Authorization: Bearer <token>` header is decoded during the WebSocket
/// handshake and the resulting claims are threaded through to every `WsClient`
/// so gateway handlers can call `client.claims()` for auth decisions.
pub fn ws_route(
    runtime: &'static GatewayRuntime,
    registry: &'static ConnectionRegistry,
    container: &'static FrozenDiContainer,
) -> MethodRouter {
    let handler = move |ws: WebSocketUpgrade, headers: HeaderMap| async move {
        // Same Bearer → signed-cookie pipeline as the HTTP boundary; the
        // handshake authenticates once and never loads a server-side session.
        let claims = crate::auth::extract::extract_claims(&headers, container);
        ws.on_upgrade(move |socket| handle_socket(socket, runtime, registry, claims))
    };
    get(handler)
}

/// Drive one upgraded socket to completion: register, pump, dispatch, drain.
async fn handle_socket(
    socket: WebSocket,
    runtime: &'static GatewayRuntime,
    registry: &'static ConnectionRegistry,
    claims: Option<Arc<Claims>>,
) {
    let (mut sink, mut stream) = socket.split();

    // Outbound queue: any task enqueues, this socket's writer drains.
    let (tx, mut rx) = mpsc::unbounded_channel::<WsMessage>();
    let id = registry.register(tx, claims.clone());
    let client = WsClient::__new(id, registry, claims);

    // One-shot signal: when the writer exits for *any* reason (peer error,
    // server-initiated Close, or channel closed), the reader is unblocked so it
    // stops polling the stream and runs on_disconnect + unregister. Without this,
    // a server-initiated close would leave the reader blocked on stream.next()
    // indefinitely if the peer never sends a Close echo.
    let (close_tx, mut close_rx) = oneshot::channel::<()>();

    // Writer half — owns the sink; exits when the queue closes or the peer dies.
    let writer = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            let frame = match msg {
                WsMessage::Text(arc) => Message::Text(String::from(arc.as_ref())),
                WsMessage::Close => {
                    // Send close frame then exit immediately — do NOT loop back
                    // to rx.recv(), which would keep the sink open indefinitely.
                    let _ = sink.send(Message::Close(None)).await;
                    break;
                }
            };
            if sink.send(frame).await.is_err() {
                break;
            }
        }
        // Dropping close_tx signals the reader regardless of why we exited.
        drop(close_tx);
    });

    (runtime.on_connect)(client.clone()).await;

    // Reader half — routes inbound frames to subscribed handlers.
    // Also watches for the writer-exit signal so a server-initiated close
    // (WsMessage::Close enqueued by a handler) terminates the reader promptly.
    loop {
        tokio::select! {
            biased;
            // Writer exited (server-initiated close or peer write error).
            _ = &mut close_rx => break,
            frame = stream.next() => match frame {
                None => break,
                Some(Err(_)) => break,
                Some(Ok(frame)) => match frame {
                    Message::Text(text) => dispatch_event(runtime, &client, &text).await,
                    Message::Binary(_) => { /* binary multiplexing not enabled */ }
                    Message::Close(_) => break,
                    Message::Ping(_) | Message::Pong(_) => { /* axum auto-replies to pings */ }
                },
            }
        }
    }

    (runtime.on_disconnect)(client.clone()).await;
    registry.unregister(id);
    writer.abort();
}

/// Parse one `{ "event": ..., "data": ... }` envelope and invoke its handler.
/// Unknown events and malformed frames are ignored (a hostile client can't
/// crash the dispatcher).
async fn dispatch_event(runtime: &'static GatewayRuntime, client: &WsClient, raw: &str) {
    let Ok(value) = serde_json::from_str::<serde_json::Value>(raw) else {
        return;
    };
    let Some(event) = value.get("event").and_then(|e| e.as_str()) else {
        return;
    };
    let Some(handler) = runtime.handler(event) else {
        return;
    };

    let data = value
        .get("data")
        .cloned()
        .unwrap_or(serde_json::Value::Null);
    let data_str: Arc<str> = Arc::from(serde_json::to_string(&data).unwrap_or_default());

    // Handler errors are swallowed at the transport edge — gateways own their
    // own error-to-client signalling (e.g. emitting an `error` event).
    let _ = handler(client.clone(), data_str).await;
}