//! Real-time WebSocket event streaming.

use crate::error::{ApiError, ApiResult};
use crate::state::AppState;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::extract::{Query, State};
use axum::http::HeaderMap;
use axum::http::header::COOKIE;
use axum::response::IntoResponse;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;

#[derive(Debug, Deserialize)]
pub struct WsAuthQuery {
    pub token: Option<String>,
    pub session: Option<String>,
}

/// Extract the NX9 session identifier from the browser session cookie.
///
/// The WebUI authenticates through the HttpOnly `nx9_session` cookie.
/// WebSocket upgrades do not pass through the normal REST authentication
/// middleware, so the cookie must be authenticated explicitly here.
fn extract_session_cookie(headers: &HeaderMap) -> Option<&str> {
    headers
        .get(COOKIE)
        .and_then(|value| value.to_str().ok())
        .and_then(|cookies| {
            cookies
                .split(';')
                .map(str::trim)
                .find_map(|cookie| cookie.strip_prefix("nx9_session="))
        })
        .map(str::trim)
        .filter(|session_id| !session_id.is_empty())
}

/// Authenticate a WebSocket request.
///
/// Authentication precedence:
///
/// 1. Explicit API token: `?token=...`
/// 2. Explicit session: `?session=...`
/// 3. Browser session cookie: `nx9_session=...`
///
/// The browser WebUI uses the HttpOnly session cookie, so no credential
/// needs to be exposed in the WebSocket URL.
async fn authenticate_websocket(
    state: &AppState,
    query: &WsAuthQuery,
    headers: &HeaderMap,
) -> bool {
    if let Some(raw_token) = query.token.as_deref() {
        return state.auth.authenticate_token(raw_token).await.is_ok();
    }

    if let Some(session_id) = query.session.as_deref() {
        return state.auth.authenticate_session(session_id).await.is_ok();
    }

    if let Some(session_id) = extract_session_cookie(headers) {
        return state.auth.authenticate_session(session_id).await.is_ok();
    }

    false
}

/// GET /api/v1/ws
pub async fn ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    Query(query): Query<WsAuthQuery>,
    headers: HeaderMap,
) -> ApiResult<impl IntoResponse> {
    if !authenticate_websocket(&state, &query, &headers).await {
        return Err(ApiError::Unauthenticated(
            "WebSocket authentication required. Supply a valid API token, session, or nx9_session cookie."
                .to_string(),
        ));
    }

    Ok(ws.on_upgrade(move |socket| handle_socket(socket, state)))
}

async fn handle_socket(socket: WebSocket, state: AppState) {
    let (mut sender, mut receiver) = socket.split();
    let mut rx = state.event_tx.subscribe();

    // Stream broadcast events to the connected WebSocket client.
    let mut send_task = tokio::spawn(async move {
        while let Ok(event) = rx.recv().await {
            if let Ok(json) = serde_json::to_string(&event) {
                let msg = Message::Text(json.into());

                if sender.send(msg).await.is_err() {
                    break;
                }
            }
        }
    });

    // Receive loop handles client close frames and keeps the connection alive.
    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(msg)) = receiver.next().await {
            if let Message::Close(_) = msg {
                break;
            }
        }
    });

    // If either side terminates, stop the other task.
    tokio::select! {
        _ = (&mut send_task) => {
            recv_task.abort();
        }
        _ = (&mut recv_task) => {
            send_task.abort();
        }
    }
}