//! WebSocket tunnel server - accepts incoming tunnel connections.
//!
//! Provides `accept_tunnel_connection` which takes an axum WebSocket,
//! sets up a TunnelEndpoint, and returns a handle for the connection.

use crate::state::LivenessToken;
use crate::tunnel::{create_endpoint, Dispatch};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, oneshot};

/// Handle to an accepted tunnel connection.
///
/// Returned by `accept_tunnel_connection`. The connection runs in the background.
/// Use `disconnected_rx` to be notified when the connection drops.
pub struct ServerConnectionHandle {
    /// RPC endpoint for calling the remote side
    pub endpoint: Arc<crate::tunnel::TunnelEndpoint>,
    /// Liveness token for this connection. Check `liveness.is_alive()` or pass to
    /// `ConnectionRegistry::insert`.
    pub liveness: LivenessToken,
    /// Notified when the connection is closed (recv returns Ok(()) or Err)
    pub disconnected_rx: oneshot::Receiver<()>,
}

/// Accept an incoming WebSocket tunnel connection.
///
/// Sets up a TunnelEndpoint with the given dispatchers and spawns the
/// connection handler. Returns immediately with a handle.
///
/// The connection is ready for RPC immediately — the remote side can
/// call any registered dispatcher, and you can call remote methods
/// through `handle.endpoint`.
pub async fn accept_tunnel_connection(
    ws: axum::extract::ws::WebSocket,
    dispatchers: Vec<Box<dyn Dispatch>>,
    rpc_timeout: Duration,
) -> ServerConnectionHandle {
    let (endpoint, ws_to_ep_tx, ep_to_ws_rx) = create_endpoint(dispatchers, rpc_timeout);
    let connected = Arc::new(AtomicBool::new(true));
    let (disconnected_tx, disconnected_rx) = oneshot::channel();

    let connected_clone = connected.clone();
    tokio::spawn(run_server_connection(
        ws,
        ws_to_ep_tx,
        ep_to_ws_rx,
        connected_clone,
        disconnected_tx,
    ));

    ServerConnectionHandle {
        endpoint,
        liveness: LivenessToken::new(connected),
        disconnected_rx,
    }
}

async fn run_server_connection(
    ws: axum::extract::ws::WebSocket,
    ws_to_ep_tx: mpsc::UnboundedSender<String>,
    mut ep_to_ws_rx: mpsc::UnboundedReceiver<String>,
    connected: Arc<AtomicBool>,
    disconnected_tx: oneshot::Sender<()>,
) {
    use axum::extract::ws::Message;
    use futures_util::{SinkExt, StreamExt};

    let (mut writer, mut reader) = ws.split();

    // Spawn reader: forward all text messages to endpoint
    let ws_to_ep_tx_clone = ws_to_ep_tx.clone();
    let mut reader_handle = tokio::spawn(async move {
        while let Some(msg) = reader.next().await {
            match msg {
                Ok(Message::Text(text)) => {
                    let _ = ws_to_ep_tx_clone.send(text.to_string());
                }
                Ok(Message::Close(_)) => break,
                Err(e) => {
                    tracing::warn!("[TUNNEL-SERVER] Read error: {}", e);
                    break;
                }
                _ => {} // Ping/Pong handled by axum
            }
        }
    });

    // Main loop: forward outgoing messages from endpoint to WebSocket
    loop {
        tokio::select! {
            msg = ep_to_ws_rx.recv() => {
                match msg {
                    Some(text) => {
                        if writer.send(Message::Text(text.into())).await.is_err() {
                            break;
                        }
                    }
                    None => break,
                }
            }
            _ = &mut reader_handle => {
                break;
            }
        }
    }

    // Connection closed
    connected.store(false, Ordering::Release);
    let _ = disconnected_tx.send(());
    drop(ws_to_ep_tx); // Signal endpoint reader_loop to stop
}