//! WebSocket interface — real-time bidirectional communication.
//!
//! Enables web clients and other WS-capable consumers to connect
//! and participate in sessions. Built on Axum's WebSocket support.

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::response::IntoResponse;
use axum::routing::get;
use axum::Router;
use futures_util::stream::StreamExt;
use futures_util::SinkExt;
use std::sync::Arc;
use tracing::info;

/// Shared application state for the WS server.
pub struct WsState {
    pub store_base_dir: std::path::PathBuf,
    pub session_id: Option<String>,
}

/// Build the WebSocket router.
pub fn build_router(state: Arc<WsState>) -> Router {
    Router::new()
        .route("/ws", get(ws_handler))
        .with_state(state)
}

/// WebSocket upgrade handler.
async fn ws_handler(
    ws: WebSocketUpgrade,
    axum::extract::State(state): axum::extract::State<Arc<WsState>>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_socket(socket, state))
}

/// Handle an established WebSocket connection.
async fn handle_socket(mut socket: WebSocket, state: Arc<WsState>) {
    // Channel for sending text messages to the WebSocket send task.
    // The receiver side runs in a spawned task that forwards each
    // string as a `Message::Text` to the client.
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();

    info!("WebSocket client connected");

    // Send a welcome message
    let welcome = serde_json::json!({
        "type": "connected",
        "session": state.session_id,
        "message": "Connected to Zesdex WebSocket server"
    });
    // axum 0.8 Message::Text wraps Utf8Bytes; convert via .into()
    let _ = socket.send(Message::Text(welcome.to_string().into())).await;

    // Split the socket into sender and receiver halves
    let (mut sender, mut receiver) = socket.split();

    // Spawn task to forward messages from channel to WebSocket sender
    let send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if sender.send(Message::Text(msg.into())).await.is_err() {
                break;
            }
        }
    });

    // Receive messages from the client
    // receiver is SplitStream<WebSocket> — use StreamExt::next()
    while let Some(Ok(msg)) = receiver.next().await {
        match msg {
            Message::Text(text) => {
                // Convert Utf8Bytes -> String for JSON serialisation
                let text_str = text.to_string();
                info!("Received WS message: {text_str}");
                // Echo back for now
                let response = serde_json::json!({
                    "type": "echo",
                    "data": text_str
                });
                let _ = tx.send(response.to_string());
            }
            Message::Close(_) => break,
            _ => {}
        }
    }

    send_task.abort();
    info!("WebSocket client disconnected");
}

/// Run the WebSocket server standalone.
pub async fn run_server(port: u16) -> anyhow::Result<()> {
    let state = Arc::new(WsState {
        store_base_dir: std::path::PathBuf::from("."),
        session_id: None,
    });
    let app = build_router(state);
    let addr = std::net::SocketAddr::from(([0, 0, 0, 0], port));
    info!("WebSocket server listening on ws://{addr}");
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;
    Ok(())
}