//! Streamable HTTP MCP server.
//!
//! One `/mcp` route, three methods:
//! - **POST** — JSON-RPC requests: `initialize` (opens a session, returns its
//!   id in `Mcp-Session-Id`), `tools/list`, `tools/call`, `resources/list`,
//!   `resources/read`, `resources/subscribe`/`unsubscribe`, and the
//!   `notifications/cancelled` notification.
//! - **GET** — opens the server→client SSE stream for a session, down which
//!   `notifications/resources/updated` and `notifications/cancelled` are
//!   pushed.
//! - **DELETE** — ends a session.
//!
//! Configuration is passed at startup via [`ServeConfig`]:
//! - `projection` — Full or Compact surface exposed to clients on this port.
//! - `grants` — which write capabilities are honored; unmatched writes return
//!   `CapabilityDenied`. `None` bypasses the check (fully-trusted context).
//! - `resources` — an optional [`ResourceRegistry`]; when set, the resource
//!   methods above are served and the `resources` capability is advertised.

use std::collections::HashSet;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;

use axum::{
    Json, Router,
    extract::State,
    http::{HeaderMap, StatusCode},
    response::{
        IntoResponse,
        sse::{Event, KeepAlive, Sse},
    },
    routing::post,
};
use serde_json::{Value, json};
use tokio::net::TcpListener;
use tokio::task::JoinHandle;
use tokio_stream::wrappers::UnboundedReceiverStream;
use tokio_stream::{Stream, StreamExt};

use crate::error::{Error, Result};
use crate::resource::{ReadResourceParams, ResourceRegistry, read_result};
use crate::tool::{SurfaceProjection, ToolRegistry};

pub mod protocol;
mod session;
pub mod stdio;

use protocol::{
    CallToolParams, CancelledParams, InitializeResult, JsonRpcRequest, JsonRpcResponse,
    MCP_PROTOCOL_VERSION, ResourcesCapability, ServerCapabilities, ServerInfo, SubscribeParams,
    ToolsCapability, codes,
};
use session::{Session, Sessions, resource_updated_notification};

/// Header carrying the MCP session id, per the Streamable HTTP transport.
const SESSION_HEADER: &str = "mcp-session-id";

pub struct ServeConfig {
    pub projection: SurfaceProjection,
    pub grants: Option<HashSet<String>>,
    /// Resource surface exposed alongside the tools. `None` means the server
    /// advertises no `resources` capability and rejects `resources/*` methods.
    pub resources: Option<ResourceRegistry>,
}

impl Default for ServeConfig {
    fn default() -> Self {
        Self {
            projection: SurfaceProjection::Full,
            grants: Some(HashSet::new()),
            resources: None,
        }
    }
}

impl ServeConfig {
    /// Full surface, honoring exactly the given write-capability ids. The
    /// common shape for a standalone binary that knows which capabilities its
    /// driving session should hold.
    pub fn granting(ids: impl IntoIterator<Item = impl Into<String>>) -> Self {
        Self {
            projection: SurfaceProjection::Full,
            grants: Some(ids.into_iter().map(Into::into).collect()),
            resources: None,
        }
    }

    #[must_use]
    pub fn with_projection(mut self, projection: SurfaceProjection) -> Self {
        self.projection = projection;
        self
    }

    /// Attach a resource registry, enabling `resources/list` and
    /// `resources/read` and advertising the `resources` capability.
    #[must_use]
    pub fn with_resources(mut self, resources: ResourceRegistry) -> Self {
        self.resources = Some(resources);
        self
    }
}

/// Handle to a running MCP server.
///
/// Dropping this does *not* stop the server — call [`BoundServer::shutdown`]
/// or abort the task explicitly.
pub struct BoundServer {
    pub addr: SocketAddr,
    handle: JoinHandle<std::io::Result<()>>,
    sessions: Sessions,
}

impl BoundServer {
    /// Full URL clients POST to, e.g. `http://127.0.0.1:54321/mcp`.
    pub fn url(&self) -> String {
        format!("http://{}/mcp", self.addr)
    }

    /// Cancel every in-flight `tools/call` across all sessions, pushing a
    /// `notifications/cancelled` to each affected client. The app-initiated
    /// path — a human taking a buffer back in deox. Returns the count
    /// cancelled. For v0 (one agent = one session) this is the whole cancel
    /// surface; [`cancel_session`](Self::cancel_session) targets one client
    /// once an app tracks session ids.
    pub fn cancel_all(&self, reason: &str) -> usize {
        self.sessions.cancel_all(reason)
    }

    /// Cancel in-flight `tools/call`s for one session by its `Mcp-Session-Id`.
    /// Returns the count cancelled (0 if the session is unknown or idle).
    pub fn cancel_session(&self, session_id: &str, reason: &str) -> usize {
        self.sessions
            .get(session_id)
            .map_or(0, |s| s.cancel_all(reason))
    }

    /// Ids of the currently-connected sessions.
    pub fn sessions(&self) -> Vec<String> {
        self.sessions.ids()
    }

    pub fn shutdown(self) {
        self.handle.abort();
    }

    /// Block until the server task finishes.
    ///
    /// A standalone MCP binary calls this after [`serve`] to keep the process
    /// alive for as long as the server is listening. The task only ends on an
    /// unrecoverable listener error (or an external abort), so under normal
    /// operation this future never resolves — it is the binary's park point.
    pub async fn wait(self) -> Result<()> {
        match self.handle.await {
            Ok(io_result) => io_result.map_err(Error::from),
            Err(join_err) if join_err.is_cancelled() => Ok(()),
            Err(join_err) => Err(Error::Transport(join_err.to_string())),
        }
    }
}

#[derive(Clone)]
struct AppState {
    registry: ToolRegistry,
    config: Arc<ServeConfig>,
    sessions: Sessions,
}

/// Start an MCP server. Returns once the listener is bound; the server runs
/// in a background tokio task.
pub async fn serve(registry: ToolRegistry, bind: &str, config: ServeConfig) -> Result<BoundServer> {
    let listener = TcpListener::bind(bind).await?;
    let addr = listener.local_addr()?;

    let sessions = Sessions::default();
    let state = AppState {
        registry,
        config: Arc::new(config),
        sessions: sessions.clone(),
    };

    let app = Router::new()
        .route(
            "/mcp",
            post(handle_rpc).get(handle_sse).delete(handle_delete),
        )
        .with_state(state);

    let handle = tokio::spawn(async move { axum::serve(listener, app).await });

    tracing::info!(%addr, "kberg server listening");
    Ok(BoundServer {
        addr,
        handle,
        sessions,
    })
}

/// Look up the session named by the `Mcp-Session-Id` header, if any.
fn session_from_headers(state: &AppState, headers: &HeaderMap) -> Option<Arc<Session>> {
    headers
        .get(SESSION_HEADER)
        .and_then(|h| h.to_str().ok())
        .and_then(|id| state.sessions.get(id))
}

async fn handle_rpc(
    State(state): State<AppState>,
    headers: HeaderMap,
    Json(req): Json<Value>,
) -> impl IntoResponse {
    let req: JsonRpcRequest = match serde_json::from_value(req) {
        Ok(r) => r,
        Err(e) => {
            return (
                StatusCode::BAD_REQUEST,
                Json(JsonRpcResponse::err(
                    Value::Null,
                    codes::PARSE_ERROR,
                    e.to_string(),
                )),
            )
                .into_response();
        }
    };

    // `initialize` opens a session; every other method rides an existing one
    // named by the header (absent for a bare client, e.g. over stdio-style use).
    let is_initialize = req.method == "initialize";
    let session = session_from_headers(&state, &headers);

    let dispatched = dispatch(
        &state.registry,
        state.config.resources.as_ref(),
        state.config.projection,
        state.config.grants.as_ref(),
        session.as_deref(),
        req,
    )
    .await;

    match dispatched {
        // Notifications (no `id`) get a 202 with no body.
        None => StatusCode::ACCEPTED.into_response(),
        Some(response) => {
            let mut http = (StatusCode::OK, Json(response)).into_response();
            if is_initialize {
                let new_session = state.sessions.create();
                if let Ok(value) = new_session.id.parse() {
                    http.headers_mut().insert(SESSION_HEADER, value);
                }
            }
            http
        }
    }
}

/// GET `/mcp` — open the server→client SSE stream for a session. This is the
/// channel `notifications/resources/updated` and `notifications/cancelled`
/// travel down. Requires a valid `Mcp-Session-Id`.
async fn handle_sse(State(state): State<AppState>, headers: HeaderMap) -> axum::response::Response {
    let Some(session) = session_from_headers(&state, &headers) else {
        return (StatusCode::BAD_REQUEST, "missing or unknown Mcp-Session-Id").into_response();
    };

    let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
    session.attach_out(tx.clone());

    // Forward resource updates this session is subscribed to. The task lives as
    // long as the SSE stream: once the client disconnects, `rx` drops and the
    // `tx.send` below fails, ending the loop.
    if let Some(resources) = &state.config.resources {
        let mut updates = resources.subscribe_updates();
        let session = session.clone();
        tokio::spawn(async move {
            loop {
                match updates.recv().await {
                    Ok(uri) => {
                        if session.is_subscribed(&uri)
                            && tx.send(resource_updated_notification(&uri)).is_err()
                        {
                            break; // client gone
                        }
                    }
                    // A slow session that lagged past the buffer drops those
                    // updates and keeps streaming; it re-reads on next use.
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {}
                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                }
            }
        });
    }

    let stream = UnboundedReceiverStream::new(rx)
        .map(|note| Ok::<Event, Infallible>(Event::default().data(note.to_string())));
    sse_response(stream)
}

fn sse_response<S>(stream: S) -> axum::response::Response
where
    S: Stream<Item = std::result::Result<Event, Infallible>> + Send + 'static,
{
    Sse::new(stream)
        .keep_alive(KeepAlive::default())
        .into_response()
}

/// DELETE `/mcp` — end a session (the client is done). Idempotent.
async fn handle_delete(State(state): State<AppState>, headers: HeaderMap) -> impl IntoResponse {
    if let Some(id) = headers.get(SESSION_HEADER).and_then(|h| h.to_str().ok()) {
        state.sessions.remove(id);
    }
    StatusCode::NO_CONTENT
}

/// Transport-agnostic JSON-RPC dispatch. Handles `initialize`, `tools/list`,
/// and `tools/call` against a registry. Returns `None` for a notification
/// (a request with no `id`), which carries no response. Shared by the HTTP
/// handler and the [`stdio`] transport.
pub(crate) async fn dispatch(
    registry: &ToolRegistry,
    resources: Option<&ResourceRegistry>,
    projection: SurfaceProjection,
    grants: Option<&HashSet<String>>,
    session: Option<&Session>,
    req: JsonRpcRequest,
) -> Option<JsonRpcResponse> {
    // Notifications (no `id`) carry no response, but some carry a side effect.
    let Some(id) = req.id.clone() else {
        if req.method == "notifications/cancelled"
            && let Some(session) = session
            && let Ok(params) =
                serde_json::from_value::<CancelledParams>(req.params.unwrap_or(Value::Null))
        {
            session.cancel_request(&params.request_id.to_string());
        }
        return None;
    };

    let response = match req.method.as_str() {
        "initialize" => JsonRpcResponse::ok(id, initialize_result(resources.is_some())),
        "tools/list" => {
            let specs = registry.specs_projected(projection);
            JsonRpcResponse::ok(id, json!({ "tools": specs }))
        }
        "tools/call" => {
            match serde_json::from_value::<CallToolParams>(req.params.unwrap_or(Value::Null)) {
                Ok(params) => {
                    let call = registry.call(&params.name, params.arguments, grants);
                    // With a session, run so a human/client can cancel mid-flight.
                    let result = match session {
                        Some(session) => session.run_cancellable(&id, call).await,
                        None => call.await,
                    };
                    match result {
                        Ok(r) => JsonRpcResponse::ok(id, serde_json::to_value(r).unwrap()),
                        Err(e) => rpc_err_from(id, e),
                    }
                }
                Err(e) => JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
            }
        }
        "resources/list" => match resources {
            Some(res) => JsonRpcResponse::ok(id, json!({ "resources": res.list() })),
            None => resources_unsupported(id),
        },
        "resources/read" => match resources {
            Some(res) => match serde_json::from_value::<ReadResourceParams>(
                req.params.unwrap_or(Value::Null),
            ) {
                Ok(params) => match res.read(&params.uri).await {
                    Ok(contents) => JsonRpcResponse::ok(id, read_result(contents)),
                    Err(e) => rpc_err_from(id, e),
                },
                Err(e) => JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
            },
            None => resources_unsupported(id),
        },
        method @ ("resources/subscribe" | "resources/unsubscribe") => {
            dispatch_subscription(method, resources, session, id, req.params)
        }
        other => JsonRpcResponse::err(
            id,
            codes::METHOD_NOT_FOUND,
            format!("unknown method: {other}"),
        ),
    };
    Some(response)
}

/// `resources/subscribe` / `resources/unsubscribe`. Both need a resource
/// registry (to validate the uri exists) and a session (to hold the
/// subscription) — the latter means these are HTTP-only.
fn dispatch_subscription(
    method: &str,
    resources: Option<&ResourceRegistry>,
    session: Option<&Session>,
    id: Value,
    params: Option<Value>,
) -> JsonRpcResponse {
    let Some(resources) = resources else {
        return resources_unsupported(id);
    };
    let Some(session) = session else {
        return JsonRpcResponse::err(
            id,
            codes::INVALID_PARAMS,
            "resource subscriptions require a session (HTTP transport)",
        );
    };
    let params: SubscribeParams = match serde_json::from_value(params.unwrap_or(Value::Null)) {
        Ok(p) => p,
        Err(e) => return JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
    };
    if method == "resources/subscribe" {
        if !resources.contains(&params.uri) {
            return rpc_err_from(id, Error::ResourceNotFound(params.uri));
        }
        session.subscribe(params.uri);
    } else {
        session.unsubscribe(&params.uri);
    }
    // MCP subscribe/unsubscribe return an empty result on success.
    JsonRpcResponse::ok(id, json!({}))
}

/// A `resources/*` method arrived at a server with no resource registry
/// configured. It is a method the server does not implement, not a bad request.
fn resources_unsupported(id: Value) -> JsonRpcResponse {
    JsonRpcResponse::err(id, codes::METHOD_NOT_FOUND, "server exposes no resources")
}

// id/e are consumed into the returned response.
#[allow(clippy::needless_pass_by_value)]
fn rpc_err_from(id: Value, e: Error) -> JsonRpcResponse {
    let (code, message) = match &e {
        Error::ToolNotFound(_) => (codes::METHOD_NOT_FOUND, e.to_string()),
        Error::ResourceNotFound(_) => (codes::RESOURCE_NOT_FOUND, e.to_string()),
        Error::Cancelled(_) => (codes::REQUEST_CANCELLED, e.to_string()),
        Error::InvalidArgs { .. } => (codes::INVALID_PARAMS, e.to_string()),
        Error::CapabilityDenied { .. } | Error::Refused { .. } => {
            (codes::INVALID_PARAMS, e.to_string())
        }
        _ => (codes::INTERNAL_ERROR, e.to_string()),
    };
    JsonRpcResponse::err(id, code, message)
}

fn initialize_result(resources_enabled: bool) -> Value {
    serde_json::to_value(InitializeResult {
        protocol_version: MCP_PROTOCOL_VERSION,
        capabilities: ServerCapabilities {
            tools: ToolsCapability {
                list_changed: false,
            },
            // Over HTTP, resources support subscribe + server-push updates.
            resources: resources_enabled.then_some(ResourcesCapability {
                subscribe: true,
                list_changed: false,
            }),
        },
        server_info: ServerInfo {
            name: "kberg".to_string(),
            version: env!("CARGO_PKG_VERSION").to_string(),
        },
    })
    .unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::resource::{ResourceContents, ResourceDescriptor};
    use crate::tool::{Tool, ToolCallResult, ToolKind};
    use async_trait::async_trait;

    struct Echo;
    #[async_trait]
    impl Tool for Echo {
        fn name(&self) -> &'static str {
            "echo"
        }
        fn description(&self) -> &'static str {
            "read echo"
        }
        fn kind(&self) -> ToolKind {
            ToolKind::Read
        }
        fn input_schema(&self) -> Value {
            json!({ "type": "object", "properties": {} })
        }
        async fn call(&self, _args: Value) -> crate::Result<ToolCallResult> {
            Ok(ToolCallResult::text("pong"))
        }
    }

    fn registry() -> ToolRegistry {
        let mut r = ToolRegistry::new();
        r.register(Echo);
        r
    }

    // id/params are consumed into the returned request (test helper).
    #[allow(clippy::needless_pass_by_value)]
    fn request(id: Option<Value>, method: &str, params: Value) -> JsonRpcRequest {
        serde_json::from_value(json!({
            "jsonrpc": "2.0",
            "id": id,
            "method": method,
            "params": params,
        }))
        .unwrap()
    }

    #[tokio::test]
    async fn dispatch_initialize_and_list_and_call() {
        let r = registry();
        let grants = HashSet::new();

        let init = dispatch(
            &r,
            None,
            SurfaceProjection::Full,
            Some(&grants),
            None,
            request(Some(json!(1)), "initialize", Value::Null),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(init).unwrap();
        assert_eq!(v["result"]["serverInfo"]["name"], "kberg");
        // No resource registry → no `resources` capability advertised.
        assert!(v["result"]["capabilities"]["resources"].is_null());

        let list = dispatch(
            &r,
            None,
            SurfaceProjection::Full,
            Some(&grants),
            None,
            request(Some(json!(2)), "tools/list", Value::Null),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(list).unwrap();
        assert_eq!(v["result"]["tools"][0]["name"], "echo");

        let call = dispatch(
            &r,
            None,
            SurfaceProjection::Full,
            Some(&grants),
            None,
            request(
                Some(json!(3)),
                "tools/call",
                json!({ "name": "echo", "arguments": {} }),
            ),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(call).unwrap();
        assert_eq!(v["result"]["content"][0]["text"], "pong");
    }

    #[tokio::test]
    async fn dispatch_notification_has_no_response() {
        let r = registry();
        let resp = dispatch(
            &r,
            None,
            SurfaceProjection::Full,
            None,
            None,
            request(None, "initialize", Value::Null),
        )
        .await;
        assert!(resp.is_none());
    }

    #[tokio::test]
    async fn dispatch_unknown_method_errors() {
        let r = registry();
        let resp = dispatch(
            &r,
            None,
            SurfaceProjection::Full,
            None,
            None,
            request(Some(json!(9)), "bogus", Value::Null),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(resp).unwrap();
        assert_eq!(v["error"]["code"], codes::METHOD_NOT_FOUND);
    }

    fn resources() -> ResourceRegistry {
        let res = ResourceRegistry::new();
        res.register(
            ResourceDescriptor::new("buffer://main.rs", "main.rs").with_mime_type("text/x-rust"),
            |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
        );
        res
    }

    #[tokio::test]
    async fn dispatch_initialize_advertises_resources_when_configured() {
        let r = registry();
        let res = resources();
        let init = dispatch(
            &r,
            Some(&res),
            SurfaceProjection::Full,
            None,
            None,
            request(Some(json!(1)), "initialize", Value::Null),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(init).unwrap();
        assert_eq!(v["result"]["capabilities"]["resources"]["subscribe"], true);
    }

    #[tokio::test]
    async fn dispatch_resources_list_and_read() {
        let r = registry();
        let res = resources();

        let list = dispatch(
            &r,
            Some(&res),
            SurfaceProjection::Full,
            None,
            None,
            request(Some(json!(2)), "resources/list", Value::Null),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(list).unwrap();
        assert_eq!(v["result"]["resources"][0]["uri"], "buffer://main.rs");
        assert_eq!(v["result"]["resources"][0]["mimeType"], "text/x-rust");

        let read = dispatch(
            &r,
            Some(&res),
            SurfaceProjection::Full,
            None,
            None,
            request(
                Some(json!(3)),
                "resources/read",
                json!({ "uri": "buffer://main.rs" }),
            ),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(read).unwrap();
        assert_eq!(v["result"]["contents"][0]["text"], "fn main() {}");
    }

    #[tokio::test]
    async fn dispatch_resources_read_unknown_uri_errors() {
        let r = registry();
        let res = resources();
        let read = dispatch(
            &r,
            Some(&res),
            SurfaceProjection::Full,
            None,
            None,
            request(
                Some(json!(4)),
                "resources/read",
                json!({ "uri": "buffer://gone.rs" }),
            ),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(read).unwrap();
        assert_eq!(v["error"]["code"], codes::RESOURCE_NOT_FOUND);
    }

    #[tokio::test]
    async fn dispatch_resources_method_without_registry_is_method_not_found() {
        let r = registry();
        let resp = dispatch(
            &r,
            None,
            SurfaceProjection::Full,
            None,
            None,
            request(Some(json!(5)), "resources/list", Value::Null),
        )
        .await
        .unwrap();
        let v = serde_json::to_value(resp).unwrap();
        assert_eq!(v["error"]["code"], codes::METHOD_NOT_FOUND);
    }
}