//! HTTP server implementation

use std::sync::Arc;
use dashmap::DashMap;
use tokio::sync::{mpsc, oneshot};
use tokio_stream::{StreamExt, wrappers::UnboundedReceiverStream};
use tokio_util::sync::CancellationToken;
use super::{HttpRuntimeContext, ServiceUrl, MCP_SESSION_ID, get_mcp_session_id};
use crate::{
    shared::MessageRegistry,
    types::{RequestId, Message, Response},
    error::{Error, ErrorCode}
};
use volga::{
    App, HttpResult, HttpRequest, di::Dc, status, ok,
    auth::{BearerTokenService, Bearer},
    http::sse::Message as SseMessage, sse,
    headers::{HeaderMap, AUTHORIZATION}
};
#[cfg(feature = "server-tls")]
use volga::tls::TlsConfig;
#[cfg(feature = "tracing")]
use crate::types::notification::fmt::LOG_REGISTRY;

pub use auth_config::{AuthConfig, DefaultClaims};
pub(crate) use auth_config::{validate_permissions, validate_roles};

pub(crate) mod auth_config;

type RequestMap = Arc<DashMap<RequestId, oneshot::Sender<Message>>>;

#[derive(Clone)]
struct RequestManager {
    pending: RequestMap,
    msg_registry: Arc<MessageRegistry>,
    sender: mpsc::Sender<Result<Message, Error>>,
}

pub(super) async fn serve(
    rt: HttpRuntimeContext,
    token: CancellationToken
) {
    let pending = Arc::new(DashMap::new());
    let registry = Arc::new(MessageRegistry::new());
    let manager = RequestManager {
        pending: pending.clone(),
        msg_registry: registry.clone(),
        sender: rt.tx,
    };
    tokio::join!(
        dispatch(pending.clone(), registry.clone(), rt.rx, token.clone()),
        handle(
            rt.url, 
            rt.auth,
            #[cfg(feature = "server-tls")]
            rt.tls_config, 
            manager, 
            token.clone())
    );
}

async fn dispatch(
    pending: RequestMap,
    msg_registry: Arc<MessageRegistry>,
    mut sender_rx: mpsc::Receiver<Message>,
    token: CancellationToken
) {
    loop {
        tokio::select! {
            biased;
            _ = token.cancelled() => break,
            Some(msg) = sender_rx.recv() => {
                if let Some((_, resp_tx)) = pending.remove(&msg.full_id()) {
                    if let Err(_e) = resp_tx.send(msg) {
                        #[cfg(feature = "tracing")]
                        tracing::error!(logger = "neva", "Failed to send response: {:?}", _e);
                        token.cancel();
                    }
                } else if let Err(_e) = msg_registry.send(msg) {
                    #[cfg(feature = "tracing")]
                    tracing::error!(logger = "neva", "Failed to send server request: {:?}", _e);
                }
            }
        }
    }
}

async fn handle(
    service_url: ServiceUrl,
    auth: Option<AuthConfig>,
    #[cfg(feature = "server-tls")]
    tls: Option<TlsConfig>,
    manager: RequestManager,
    token: CancellationToken
) {
    let root = "/";
    let mut server = App::new()
        .bind(service_url.addr)
        .with_no_delay()
        .without_greeter();
    
    if let Some(auth) = auth {
        let (auth, rules) = auth.into_parts();
        server = server.with_bearer_auth(|_| auth);
        server.authorize(rules);
    }

    #[cfg(feature = "server-tls")]
    if let Some(tls) = tls {
        server = server.set_tls(tls);
    }
    
    server
        .add_singleton(manager)
        .map_err(handle_http_error)
        .group(service_url.endpoint, |mcp| { 
            mcp.map_get(root, handle_connection);
            mcp.map_post(root, handle_message);
            mcp.map_delete(root, handle_session_end); 
        });
    
    if let Err(_e) = server.run().await {
        #[cfg(feature = "tracing")]
        tracing::error!(logger = "neva", "HTTP Server was shutdown: {:?}", _e);
        token.cancel();
    }
}

async fn handle_session_end(req: HttpRequest) -> HttpResult {
    let Some(id) = get_mcp_session_id(req.headers()) else {
        return status!(400);
    };

    let manager: Dc<RequestManager> = req.extract()?;
    
    #[cfg(feature = "tracing")]
    LOG_REGISTRY.unregister(&id);
    manager.msg_registry.unregister(&id);
    
    ok!([(MCP_SESSION_ID, id.to_string())])
}

async fn handle_connection(req: HttpRequest) -> HttpResult {
    let Some(id) = get_mcp_session_id(req.headers()) else { 
        return status!(400);
    };

    let manager: Dc<RequestManager> = req.extract()?;
    
    let (_log_tx, log_rx) = mpsc::unbounded_channel::<Message>();
    let (msg_tx, msg_rx) = mpsc::unbounded_channel::<Message>();
    
    #[cfg(feature = "tracing")]
    LOG_REGISTRY.register(id, _log_tx);
    manager.msg_registry.register(id, msg_tx);
    
    let stream = futures_util::stream::select(
        UnboundedReceiverStream::new(log_rx), 
        UnboundedReceiverStream::new(msg_rx));
    
    sse!(stream.map(handle_sse_message); [
        (MCP_SESSION_ID, id.to_string())
    ])
}

async fn handle_message(req: HttpRequest) -> HttpResult {
    let bts: Option<BearerTokenService> = req.extract()?;
    let manager: Dc<RequestManager> = req.extract()?;

    let mut headers = req.headers().clone();
    let id = get_or_create_mcp_session(&headers);

    // JSON-RPC 2.0 §5.1 / §6: decoding failures must be returned as a JSON-RPC
    // error object, not as a transport-level error. read_message distinguishes
    // the two cases required by the spec.
    let msg = match read_message(req).await {
        Ok(msg) => msg,
        Err(code) => {
            let resp = Response::error(RequestId::Null, Error::from(code));
            return ok!(resp; [(MCP_SESSION_ID, id.to_string())]);
        }
    };
    if let Message::Notification(_) = msg {
        return status!(202; [
            (MCP_SESSION_ID, id.to_string())
        ]);
    }

    // A batch whose items are all notifications/responses produces no reply
    // per JSON-RPC 2.0 #6. Return 202 immediately without allocating a
    // pending entry — otherwise the oneshot receiver would hang forever.
    // session_id must still be set so Response envelopes inside the batch
    // resolve pending entries by the full session_id/request_id key.
    if let Message::Batch(ref batch) = msg && !batch.has_requests() && !batch.has_error_responses() {
        let msg = msg.set_session_id(id);
        manager.sender.send(Ok(msg)).await.map_err(sender_error)?;
        return status!(202; [
            (MCP_SESSION_ID, id.to_string())
        ]);
    }

    let claims = bts
        .and_then(|bts| {
            headers.get(AUTHORIZATION)
                .and_then(|bearer| Bearer::try_from(bearer).ok())
                .and_then(|bearer| bts.decode::<DefaultClaims>(bearer).ok())
        })
        .unwrap_or_default();
    
    headers.remove(AUTHORIZATION);
    
    let msg = msg
        .set_session_id(id)
        .set_claims(claims)
        .set_headers(headers);
    
    let (resp_tx, resp_rx) = oneshot::channel::<Message>();
    manager.pending.insert(msg.full_id(), resp_tx);
    manager.sender.send(Ok(msg))
        .await
        .map_err(sender_error)?;
    let resp = resp_rx
        .await
        .map_err(receiver_error)?;
    
    ok!(resp; [
        (MCP_SESSION_ID, id.to_string())
    ])
}

/// Reads and decodes a JSON-RPC [`Message`] from the HTTP request body.
///
/// Returns [`ErrorCode`] on failure so the caller can build a spec-compliant
/// error response without promoting the failure to an HTTP-level error:
/// - [`ErrorCode::ParseError`] (`-32700`) — body is not syntactically valid JSON,
///   or the underlying transport stream fails mid-read.
/// - [`ErrorCode::InvalidRequest`] (`-32600`) — body is valid JSON but not a
///   valid JSON-RPC message (e.g. an empty batch `[]` or an unrecognised shape).
#[inline]
async fn read_message(req: HttpRequest) -> Result<Message, ErrorCode> {
    let mut body_data_stream = req.into_body().into_data_stream();
    let mut buf = bytes::BytesMut::new();

    while let Some(chunk) = body_data_stream.next().await {
        let chunk = chunk.map_err(|_| ErrorCode::ParseError)?;
        buf.extend_from_slice(&chunk);
    }

    // Two-step decode per JSON-RPC 2.0 §5.1:
    //   1. Validate JSON syntax → ParseError on failure.
    //   2. Validate message shape → InvalidRequest on failure.
    let value: serde_json::Value = serde_json::from_slice(&buf)
        .map_err(|_| ErrorCode::ParseError)?;

    serde_json::from_value::<Message>(value)
        .map_err(|_| ErrorCode::InvalidRequest)
}

async fn handle_http_error(_err: volga::error::Error) {
    #[cfg(feature = "tracing")]
    tracing::error!(logger = "neva", "HTTP error: {:?}", _err)
}

fn sender_error(err: mpsc::error::SendError<Result<Message, Error>>) -> volga::error::Error {
    volga::error::Error::new("/", err.to_string())
}

fn receiver_error(err: oneshot::error::RecvError) -> volga::error::Error {
    volga::error::Error::new("/", err.to_string())
}

fn handle_sse_message(msg: Message) -> Result<SseMessage, volga::error::Error> {
    Ok(SseMessage::new()
        .id(uuid::Uuid::new_v4().to_string())
        .json(msg))
}

/// Fetches the [`MCP_SESSION_ID`] header value or created the new [`uuid::Uuid`] if `None`
#[inline]
fn get_or_create_mcp_session(headers: &HeaderMap) -> uuid::Uuid {
    get_mcp_session_id(headers).unwrap_or_else(uuid::Uuid::new_v4)
}
