//! End-to-end tests for the Streamable HTTP MCP server: drive a real bound
//! server over HTTP with a plain `reqwest` client (no `Agent`), exercising
//! `initialize`, `tools/list` (both projections), and `tools/call` across the
//! read / write-granted / write-denied / refusal / not-found paths.

#![cfg(feature = "server")]

use std::collections::HashSet;
use std::time::Duration;

use async_trait::async_trait;
use kberg::server::{self, ServeConfig};
use kberg::{
    Refusal, ResourceContents, ResourceDescriptor, ResourceRegistry, Result, SurfaceProjection,
    Tool, ToolCallResult, ToolKind, ToolRegistry, WriteCapability,
};
use serde_json::{Value, json};

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

struct Save;
#[async_trait]
impl Tool for Save {
    fn name(&self) -> &'static str {
        "save"
    }
    fn description(&self) -> &'static str {
        "write, gated on `demo.save`"
    }
    fn kind(&self) -> ToolKind {
        ToolKind::Write(WriteCapability::new("demo.save", "persist a value"))
    }
    fn input_schema(&self) -> Value {
        json!({ "type": "object", "properties": { "v": { "type": "string" } } })
    }
    async fn call(&self, args: Value) -> Result<ToolCallResult> {
        let v = args.get("v").and_then(Value::as_str).unwrap_or("");
        Ok(ToolCallResult::text(format!("saved {v}")))
    }
}

fn registry() -> ToolRegistry {
    let mut r = ToolRegistry::new();
    r.register(Ping);
    r.register(Save);
    r.register(Refusal::new("nuke", "not offered", "never exposed"));
    r
}

/// A read tool that never returns on its own — only cancellation ends it.
struct Hang;
#[async_trait]
impl Tool for Hang {
    fn name(&self) -> &'static str {
        "hang"
    }
    fn description(&self) -> &'static str {
        "read-only; blocks until cancelled"
    }
    fn kind(&self) -> ToolKind {
        ToolKind::Read
    }
    fn input_schema(&self) -> Value {
        json!({ "type": "object", "properties": {} })
    }
    async fn call(&self, _args: Value) -> Result<ToolCallResult> {
        std::future::pending::<()>().await;
        unreachable!()
    }
}

async fn rpc(client: &reqwest::Client, url: &str, method: &str, params: Value) -> Value {
    client
        .post(url)
        .json(&json!({ "jsonrpc": "2.0", "id": 1, "method": method, "params": params }))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap()
}

/// POST a request under a session id and return the parsed JSON-RPC response.
async fn rpc_in_session(
    client: &reqwest::Client,
    url: &str,
    session: &str,
    id: i64,
    method: &str,
    params: Value,
) -> Value {
    client
        .post(url)
        .header("Mcp-Session-Id", session)
        .json(&json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap()
}

/// Initialize and return the server-assigned `Mcp-Session-Id`.
async fn open_session(client: &reqwest::Client, url: &str) -> String {
    let resp = client
        .post(url)
        .json(&json!({ "jsonrpc": "2.0", "id": 1, "method": "initialize" }))
        .send()
        .await
        .unwrap();
    resp.headers()
        .get("Mcp-Session-Id")
        .expect("initialize returns a session id")
        .to_str()
        .unwrap()
        .to_string()
}

/// Read the next `data:` JSON payload off an SSE response, buffering across
/// chunks. Returns `None` at end of stream.
async fn next_sse_json(resp: &mut reqwest::Response, buf: &mut String) -> Option<Value> {
    loop {
        if let Some(pos) = buf.find("data:")
            && let Some(nl) = buf[pos..].find('\n')
        {
            let line = buf[pos + 5..pos + nl].trim().to_string();
            *buf = buf[pos + nl + 1..].to_string();
            if !line.is_empty()
                && let Ok(v) = serde_json::from_str::<Value>(&line)
            {
                return Some(v);
            }
            continue;
        }
        match resp.chunk().await.ok()? {
            Some(bytes) => buf.push_str(&String::from_utf8_lossy(&bytes)),
            None => return None,
        }
    }
}

#[tokio::test]
async fn full_surface_round_trip() {
    let bound = server::serve(
        registry(),
        "127.0.0.1:0",
        ServeConfig::granting(["demo.save"]),
    )
    .await
    .unwrap();
    let url = bound.url();
    let client = reqwest::Client::new();

    // initialize
    let init = rpc(&client, &url, "initialize", Value::Null).await;
    assert_eq!(init["result"]["serverInfo"]["name"], "kberg");
    assert!(init["result"]["protocolVersion"].is_string());

    // tools/list — full surface has all three, with kind/refusal extensions
    let list = rpc(&client, &url, "tools/list", Value::Null).await;
    let tools = list["result"]["tools"].as_array().unwrap();
    assert_eq!(tools.len(), 3);
    let save = tools.iter().find(|t| t["name"] == "save").unwrap();
    assert_eq!(save["kburgKind"], "write");
    assert_eq!(save["kburgCapability"]["id"], "demo.save");
    let nuke = tools.iter().find(|t| t["name"] == "nuke").unwrap();
    assert!(nuke["kburgRefusalReason"].is_string());

    // tools/call — read works
    let ping = rpc(
        &client,
        &url,
        "tools/call",
        json!({ "name": "ping", "arguments": {} }),
    )
    .await;
    assert_eq!(ping["result"]["content"][0]["text"], "pong");

    // tools/call — granted write works
    let save = rpc(
        &client,
        &url,
        "tools/call",
        json!({ "name": "save", "arguments": { "v": "x" } }),
    )
    .await;
    assert_eq!(save["result"]["content"][0]["text"], "saved x");

    // tools/call — refusal returns a JSON-RPC error, not a tool result
    let nuke = rpc(
        &client,
        &url,
        "tools/call",
        json!({ "name": "nuke", "arguments": {} }),
    )
    .await;
    assert!(nuke["error"].is_object());

    // tools/call — unknown tool errors
    let miss = rpc(
        &client,
        &url,
        "tools/call",
        json!({ "name": "ghost", "arguments": {} }),
    )
    .await;
    assert!(miss["error"].is_object());

    bound.shutdown();
}

#[tokio::test]
async fn ungranted_write_is_denied_over_the_wire() {
    // Server started with an empty grant set: `save` must be refused.
    let bound = server::serve(
        registry(),
        "127.0.0.1:0",
        ServeConfig::granting(Vec::<String>::new()),
    )
    .await
    .unwrap();
    let url = bound.url();
    let client = reqwest::Client::new();

    let save = rpc(
        &client,
        &url,
        "tools/call",
        json!({ "name": "save", "arguments": { "v": "x" } }),
    )
    .await;
    let msg = save["error"]["message"].as_str().unwrap();
    assert!(
        msg.contains("demo.save"),
        "denial names the capability: {msg}"
    );
    assert!(
        msg.contains("do not retry"),
        "denial tells small models not to retry: {msg}"
    );

    bound.shutdown();
}

#[tokio::test]
async fn resources_listed_read_and_removed_over_the_wire() {
    // A resource registry the app mutates at runtime, mirroring deox opening
    // and closing buffers while the server stays up.
    let resources = ResourceRegistry::new();
    resources.register(
        ResourceDescriptor::new("buffer://main.rs", "main.rs").with_mime_type("text/x-rust"),
        |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
    );

    let bound = server::serve(
        registry(),
        "127.0.0.1:0",
        ServeConfig::granting(Vec::<String>::new()).with_resources(resources.clone()),
    )
    .await
    .unwrap();
    let url = bound.url();
    let client = reqwest::Client::new();

    // initialize advertises the resources capability with subscribe support.
    let init = rpc(&client, &url, "initialize", Value::Null).await;
    assert_eq!(
        init["result"]["capabilities"]["resources"]["subscribe"],
        true
    );

    // resources/list
    let list = rpc(&client, &url, "resources/list", Value::Null).await;
    let items = list["result"]["resources"].as_array().unwrap();
    assert_eq!(items.len(), 1);
    assert_eq!(items[0]["uri"], "buffer://main.rs");
    assert_eq!(items[0]["mimeType"], "text/x-rust");

    // resources/read
    let read = rpc(
        &client,
        &url,
        "resources/read",
        json!({ "uri": "buffer://main.rs" }),
    )
    .await;
    assert_eq!(read["result"]["contents"][0]["text"], "fn main() {}");

    // Close the buffer at runtime; the running server observes it (shared Arc).
    assert!(resources.remove("buffer://main.rs"));
    let gone = rpc(
        &client,
        &url,
        "resources/read",
        json!({ "uri": "buffer://main.rs" }),
    )
    .await;
    assert_eq!(gone["error"]["code"], -32002); // RESOURCE_NOT_FOUND

    bound.shutdown();
}

#[tokio::test]
async fn compact_projection_hides_write_tools() {
    let bound = server::serve(
        registry(),
        "127.0.0.1:0",
        ServeConfig::granting(["demo.save"]).with_projection(SurfaceProjection::Compact),
    )
    .await
    .unwrap();
    let url = bound.url();
    let client = reqwest::Client::new();

    let list = rpc(&client, &url, "tools/list", Value::Null).await;
    let names: HashSet<String> = list["result"]["tools"]
        .as_array()
        .unwrap()
        .iter()
        .map(|t| t["name"].as_str().unwrap().to_string())
        .collect();
    assert!(names.contains("ping"));
    assert!(!names.contains("save"));
    assert!(!names.contains("nuke"));

    bound.shutdown();
}

#[tokio::test]
async fn resource_update_is_pushed_over_sse() {
    // The deox flow: agent subscribes to a buffer, the human edits it, the
    // agent's SSE stream receives a `resources/updated` push.
    let resources = ResourceRegistry::new();
    resources.register(
        ResourceDescriptor::new("buffer://main.rs", "main.rs"),
        |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
    );
    let bound = server::serve(
        registry(),
        "127.0.0.1:0",
        ServeConfig::granting(Vec::<String>::new()).with_resources(resources.clone()),
    )
    .await
    .unwrap();
    let url = bound.url();
    let client = reqwest::Client::new();

    let session = open_session(&client, &url).await;

    // Open the server->client SSE stream (registers the update forwarder).
    let mut sse = client
        .get(&url)
        .header("Mcp-Session-Id", &session)
        .send()
        .await
        .unwrap();
    let mut buf = String::new();

    // Subscribe to the buffer.
    let sub = rpc_in_session(
        &client,
        &url,
        &session,
        2,
        "resources/subscribe",
        json!({ "uri": "buffer://main.rs" }),
    )
    .await;
    assert!(
        sub["result"].is_object(),
        "subscribe returns an empty result: {sub}"
    );

    // The app mutates the buffer and announces it.
    resources.notify_updated("buffer://main.rs");

    let note = tokio::time::timeout(Duration::from_secs(2), next_sse_json(&mut sse, &mut buf))
        .await
        .expect("an update arrives before the timeout")
        .expect("the stream yields a notification");
    assert_eq!(note["method"], "notifications/resources/updated");
    assert_eq!(note["params"]["uri"], "buffer://main.rs");

    bound.shutdown();
}

#[tokio::test]
async fn app_cancel_aborts_call_and_notifies_over_sse() {
    // The deox cancel: a call is in flight, the human cancels, the server
    // aborts the call and tells the agent over SSE.
    let mut reg = ToolRegistry::new();
    reg.register(Hang);
    let bound = server::serve(
        reg,
        "127.0.0.1:0",
        ServeConfig::granting(Vec::<String>::new()),
    )
    .await
    .unwrap();
    let url = bound.url();
    let client = reqwest::Client::new();

    let session = open_session(&client, &url).await;
    let mut sse = client
        .get(&url)
        .header("Mcp-Session-Id", &session)
        .send()
        .await
        .unwrap();
    let mut buf = String::new();

    // Fire a call that hangs until cancelled.
    let call = {
        let client = client.clone();
        let url = url.clone();
        let session = session.clone();
        tokio::spawn(async move {
            rpc_in_session(
                &client,
                &url,
                &session,
                7,
                "tools/call",
                json!({ "name": "hang", "arguments": {} }),
            )
            .await
        })
    };

    // Let it register as in-flight, then cancel it app-side.
    tokio::time::sleep(Duration::from_millis(200)).await;
    assert_eq!(bound.cancel_all("human took the buffer back"), 1);

    // The agent's SSE stream receives a cancelled notification echoing id 7.
    let note = tokio::time::timeout(Duration::from_secs(2), next_sse_json(&mut sse, &mut buf))
        .await
        .expect("a cancelled notification arrives")
        .expect("the stream yields a notification");
    assert_eq!(note["method"], "notifications/cancelled");
    assert_eq!(note["params"]["requestId"], 7);

    // The blocked POST completes with a request-cancelled JSON-RPC error.
    let resp = tokio::time::timeout(Duration::from_secs(2), call)
        .await
        .expect("the call returns after cancellation")
        .unwrap();
    assert_eq!(resp["error"]["code"], -32800);

    bound.shutdown();
}