//! Aggregated MCP server that merges the graph, browser, proxy, and plugin tool surfaces.
//!
//! # Tool namespace
//!
//! | Prefix | Sub-server | Example |
//! | -------- | ----------- | ------- |
//! | `graph_` | `stygian-graph` | `graph_scrape`, `graph_pipeline_run` |
//! | `browser_` | `stygian-browser` | `browser_acquire`, `browser_navigate` |
//! | `proxy_` | `stygian-proxy` | `proxy_add`, `proxy_acquire` |
//! | `plugin_` | `stygian-plugin` | `plugin_apply_template`, `plugin_list_templates` |
//! | *(none)* | Aggregator cross-crate | `scrape_proxied`, `browser_proxied` |
//!
//! The aggregator strips the `graph_` prefix before forwarding calls to the
//! graph sub-server (which internally uses un-prefixed names like `scrape`).
//! Browser, proxy, and plugin tools receive their full prefixed names in their respective servers.
//!
//! Proxy guidance: proxy-backed tools are opt-in and should not be used by
//! default. Prefer direct tools (`graph_scrape`, `browser_acquire` +
//! `browser_navigate`) unless proxy routing is explicitly requested or proxy
//! use is required for a confirmed network-blocking scenario and proxies are
//! available in the pool.

use std::sync::Arc;

use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};

use stygian_browser::{BrowserPool, mcp::McpBrowserServer};
use stygian_graph::adapters::fallback::{
    FallbackChainService, default_fallback_breaker, default_primary_breaker,
};
use stygian_graph::adapters::http::{HttpAdapter, HttpConfig};
use stygian_graph::mcp::McpGraphServer;
use stygian_graph::ports::{ScrapingService, ServiceInput};
use stygian_plugin::adapters::{ExtractionEngine, PluginExtractionAdapter};
use stygian_plugin::storage::{FileTemplateStore, MemoryIdempotencyStore};
use stygian_proxy::mcp::McpProxyServer;

const SUPPORTED_PROTOCOL_VERSIONS: &[&str] = &["2025-11-25", "2025-06-18", "2024-11-05"];

// ─── Aggregator ───────────────────────────────────────────────────────────────

/// Aggregated MCP server exposing graph, browser, and proxy capabilities.
///
/// # Example
///
/// ```no_run
/// use stygian_mcp::aggregator::McpAggregator;
///
/// # async fn run() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
/// let aggregator = McpAggregator::try_new().await?;
/// aggregator.run().await
/// # }
/// ```
pub struct McpAggregator {
    browser: Arc<McpBrowserServer>,
    proxy: Arc<McpProxyServer>,
    plugin: Arc<stygian_plugin::mcp::McpPluginServer>,
    /// HTTP-first fallback chain: tries a plain HTTP scrape first, automatically
    /// falls back to plugin template extraction on failure or circuit-open.
    fallback_chain: Arc<FallbackChainService>,
    /// Cancellation token for the proxy background health-check and session-purge tasks.
    proxy_token: CancellationToken,
    /// Join handle for the proxy background tasks (health-check + session purge).
    /// Wrapped in `Option` so it can be taken and awaited in `run()` while
    /// `Drop` can abort it on any early-exit path.
    proxy_bg: Option<JoinHandle<()>>,
}

impl McpAggregator {
    /// Create the aggregator with a default browser pool and proxy manager.
    ///
    /// # Errors
    ///
    /// Returns an error if the browser pool or proxy manager cannot be
    /// initialised (e.g. required system binaries are missing).
    pub async fn try_new() -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let pool = BrowserPool::new(stygian_browser::BrowserConfig::default()).await?;
        let browser = Arc::new(McpBrowserServer::new(pool));
        let proxy = Arc::new(McpProxyServer::new()?);
        let (proxy_token, proxy_bg) = proxy.start_background();

        // Shared template and idempotency stores — used by both the plugin MCP server
        // and the plugin extraction adapter inside the fallback chain, so that templates
        // created via plugin_create_template are immediately usable in fallback pipelines.
        let templates_dir = std::path::PathBuf::from("./plugin-templates");
        let template_store: Arc<dyn stygian_plugin::ports::PluginTemplateStore> =
            Arc::new(FileTemplateStore::new(templates_dir));
        let idempotency_store: Arc<dyn stygian_plugin::ports::IdempotencyKeyStore> =
            Arc::new(MemoryIdempotencyStore::new());

        let plugin = Arc::new(stygian_plugin::mcp::McpPluginServer::with_adapters(
            Arc::clone(&template_store),
            Arc::new(ExtractionEngine),
            Arc::clone(&idempotency_store),
        ));

        // Build fallback chain: plain HTTP (primary) → plugin extraction (fallback).
        // The HTTP adapter uses default anti-bot configuration; the plugin adapter uses
        // the same stores as the MCP server so templates are immediately available.
        let http_primary = Arc::new(HttpAdapter::with_config(HttpConfig::default()));
        let plugin_fallback = Arc::new(PluginExtractionAdapter::new(
            template_store,
            Arc::new(ExtractionEngine),
            idempotency_store,
        ));
        let fallback_chain = Arc::new(
            FallbackChainService::builder()
                .add(http_primary, default_primary_breaker())
                .add(plugin_fallback, default_fallback_breaker())
                .named("http-to-plugin")
                .build(),
        );
        info!(
            "fallback chain '{}' built with {} services",
            fallback_chain.name(),
            fallback_chain.len()
        );

        Ok(Self {
            browser,
            proxy,
            plugin,
            fallback_chain,
            proxy_token,
            proxy_bg: Some(proxy_bg),
        })
    }

    /// Run the aggregated MCP server over stdin/stdout JSON-RPC 2.0.
    ///
    /// Reads newline-delimited JSON requests and writes newline-delimited
    /// JSON responses.  Runs until stdin reaches EOF.
    ///
    /// # Errors
    ///
    /// Returns an I/O error if stdin/stdout cannot be read or written.
    pub async fn run(mut self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        info!("stygian-mcp aggregator starting (stdin/stdout mode)");

        let stdin = tokio::io::stdin();
        let stdout = tokio::io::stdout();
        let mut reader = BufReader::new(stdin).lines();
        let mut stdout = stdout;

        while let Some(line) = reader.next_line().await? {
            let line = line.trim().to_string();
            if line.is_empty() {
                continue;
            }
            debug!(?line, "MCP request");

            let response = match serde_json::from_str::<Value>(&line) {
                Ok(req) => self.handle(&req).await,
                Err(e) => Some(error_response(
                    &Value::Null,
                    -32700,
                    &format!("Parse error: {e}"),
                )),
            };

            // JSON-RPC notifications must not produce responses.
            if let Some(response) = response {
                let id = response.get("id").cloned().unwrap_or(Value::Null);
                let mut out = serde_json::to_string(&response).unwrap_or_else(|e| {
                    tracing::error!("BUG: failed to serialize MCP response: {e}");
                    serde_json::to_string(&error_response(
                        &id,
                        -32603,
                        "Internal serialization error",
                    ))
                    .unwrap_or_default()
                });
                out.push('\n');
                stdout.write_all(out.as_bytes()).await?;
                stdout.flush().await?;
            }
        }

        info!("stygian-mcp aggregator stopping (stdin closed)");

        // Shut down proxy background tasks (health-check + session purge).
        self.proxy_token.cancel();
        if let Some(bg) = self.proxy_bg.take()
            && let Err(e) = bg.await
        {
            warn!("proxy background task panicked during shutdown: {e:?}");
        }

        Ok(())
    }

    // ── Internal dispatch ─────────────────────────────────────────────────────

    async fn handle(&self, req: &Value) -> Option<Value> {
        let is_well_formed_notification = is_jsonrpc_notification(req);
        let id = req.get("id").unwrap_or(&Value::Null);

        if req.get("jsonrpc").and_then(Value::as_str) != Some("2.0") {
            return Some(error_response(
                id,
                -32600,
                "Invalid request: expected jsonrpc='2.0'",
            ));
        }

        let Some(method) = req.get("method").and_then(Value::as_str) else {
            return Some(error_response(
                id,
                -32600,
                "Invalid request: missing string 'method'",
            ));
        };

        let response = match method {
            "initialize" => handle_initialize(req),
            "initialized" | "notifications/initialized" | "ping" => ok_response(id, &json!({})),
            "tools/list" => self.handle_tools_list(id).await,
            "tools/call" => self.handle_tools_call(id, req).await,
            "resources/list" => self.handle_resources_list(id).await,
            "resources/read" => self.handle_resources_read(id, req).await,
            other => error_response(id, -32601, &format!("Method not found: {other}")),
        };

        if is_well_formed_notification {
            None
        } else {
            Some(response)
        }
    }

    // ── tools/list ────────────────────────────────────────────────────────────

    async fn handle_tools_list(&self, id: &Value) -> Value {
        let list_req = json!({"jsonrpc":"2.0","id":0,"method":"tools/list","params":{}});

        // Graph tools — prefix each name with `graph_`.
        let graph_resp = McpGraphServer::handle_request(&list_req).await;
        let graph_tools: Vec<Value> = graph_resp
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(serde_json::Value::as_array)
            .cloned()
            .unwrap_or_default()
            .into_iter()
            .map(|mut t| {
                if let Some(obj) = t.as_object_mut() {
                    let name = obj
                        .get("name")
                        .and_then(serde_json::Value::as_str)
                        .map(str::to_owned);
                    if let Some(name) = name {
                        let prefixed = format!("graph_{name}");
                        let desc = obj
                            .get("description")
                            .and_then(serde_json::Value::as_str)
                            .unwrap_or("")
                            .to_string();
                        obj.insert("name".to_string(), json!(prefixed));
                        // Prefix description so the LLM understands the namespace.
                        obj.insert("description".to_string(), json!(format!("[graph] {desc}")));
                    }
                }
                t
            })
            .collect();

        // Browser tools — already prefixed (`browser_*`).
        let browser_resp = self.browser.dispatch(&list_req).await;
        let browser_tools: Vec<Value> = browser_resp
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(serde_json::Value::as_array)
            .cloned()
            .unwrap_or_default();

        // Proxy tools — already prefixed (`proxy_*`).
        let proxy_resp = self.proxy.handle_request(&list_req).await;
        let proxy_tools: Vec<Value> = proxy_resp
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(serde_json::Value::as_array)
            .cloned()
            .unwrap_or_default();

        // Cross-crate tools implemented by the aggregator itself.
        let cross_tools = vec![
            json!({
                "name": "scrape_proxied",
                "description": "Fetch a URL through a proxy automatically selected from the pool. Acquires a proxy, performs an HTTP scrape, then releases the proxy. Returns the scraped content. IMPORTANT: Only use this tool when the user has explicitly requested proxy routing, or when IP-based rate limiting/blocking is the confirmed problem and proxies are known to be available. Do NOT use this tool by default — use graph_scrape instead. Requires at least one proxy in the pool (check proxy_pool_stats first); returns an error if the pool is empty.",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "url": { "type": "string", "description": "Target URL to scrape" },
                        "timeout_secs": { "type": "integer", "description": "Request timeout in seconds (default: 30)" }
                    },
                    "required": ["url"]
                }
            }),
            json!({
                "name": "browser_proxied",
                "description": "Navigate to a URL in a full headless browser session routed through a proxy automatically selected from the pool. Acquires a proxy and browser session, navigates, captures HTML content, then releases both. Returns navigation metadata and page HTML. IMPORTANT: Only use this tool when the user has explicitly requested proxy routing, or when both browser rendering AND IP rotation are specifically required. Do NOT use this tool by default — use browser_acquire + browser_navigate instead. Requires at least one proxy in the pool (check proxy_pool_stats first); returns an error if the pool is empty.",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "url": { "type": "string", "description": "Target URL to visit" }
                    },
                    "required": ["url"]
                }
            }),
            json!({
                "name": "scrape_with_plugin_fallback",
                "description": "Scrape a URL using a circuit-breaker-protected fallback chain: (1) Plain HTTP scrape first — fast, lightweight, no template needed. (2) If the HTTP scrape fails or the circuit opens due to repeated failures, automatically falls back to plugin template extraction using the specified template. Use this tool when you have a saved extraction template and want resilient, schema-structured results even if the site's anti-bot protection blocks plain HTTP. The `template_id` is required so the plugin fallback knows how to extract structured data; use plugin_list_templates to find available IDs. Optionally, provide `html` to skip HTTP scraping and go directly to plugin extraction.",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "url": { "type": "string", "description": "Target URL to scrape" },
                        "template_id": { "type": "string", "description": "UUID of the extraction template to apply if HTTP scraping fails" },
                        "html": { "type": "string", "description": "Optional pre-fetched HTML content. If provided, skips HTTP scraping and uses plugin extraction directly on this HTML." },
                        "idempotency_key": { "type": "string", "description": "Optional idempotency key to deduplicate repeated calls (ULID string format)" }
                    },
                    "required": ["url", "template_id"]
                }
            }),
        ];

        // Plugin tools — provided by stygian-plugin.
        let plugin_tools = self.plugin.tools_list();

        let all_tools: Vec<Value> = [
            graph_tools,
            browser_tools,
            proxy_tools,
            plugin_tools,
            cross_tools,
        ]
        .into_iter()
        .flatten()
        .collect();

        ok_response(id, &json!({ "tools": all_tools }))
    }

    // ── tools/call ────────────────────────────────────────────────────────────

    async fn handle_tools_call(&self, id: &Value, req: &Value) -> Value {
        let Some(params) = req.get("params") else {
            return error_response(id, -32602, "Missing tool 'name'");
        };
        let Some(name) = params.get("name").and_then(Value::as_str) else {
            return error_response(id, -32602, "Missing tool 'name'");
        };
        let empty = Value::Null;
        let args = params.get("arguments").unwrap_or(&empty);

        if let Some(short) = name.strip_prefix("graph_") {
            // Route to graph sub-server with un-prefixed name.
            let sub = json!({
                "jsonrpc": "2.0",
                "id": id,
                "method": "tools/call",
                "params": { "name": short, "arguments": args }
            });
            McpGraphServer::handle_request(&sub).await
        } else if name.starts_with("browser_") {
            let sub = json!({
                "jsonrpc": "2.0",
                "id": id,
                "method": "tools/call",
                "params": { "name": name, "arguments": args }
            });
            let response = self.browser.dispatch(&sub).await;
            if name == "browser_verify_stealth" {
                normalize_browser_verify_stealth_response(response)
            } else {
                response
            }
        } else if name.starts_with("proxy_") {
            let sub = json!({
                "jsonrpc": "2.0",
                "id": id,
                "method": "tools/call",
                "params": { "name": name, "arguments": args }
            });
            self.proxy.handle_request(&sub).await
        } else if name.starts_with("plugin_") {
            // Route to plugin server with full prefixed name (plugin server expects "plugin_*" names)
            let result = self.plugin.handle_tool_call(name, args).await;
            ok_response(id, &result)
        } else if name == "scrape_proxied" {
            self.tool_scrape_proxied(id, args).await
        } else if name == "browser_proxied" {
            self.tool_browser_proxied(id, args).await
        } else if name == "scrape_with_plugin_fallback" {
            self.tool_scrape_with_plugin_fallback(id, args).await
        } else {
            error_response(id, -32602, &format!("Unknown tool: {name}"))
        }
    }

    // ── resources/list ────────────────────────────────────────────────────────

    async fn handle_resources_list(&self, id: &Value) -> Value {
        let list_req = json!({"jsonrpc":"2.0","id":0,"method":"resources/list","params":{}});

        // Collect browser resources (active sessions).
        let browser_resp = self.browser.dispatch(&list_req).await;
        let browser_resources: Vec<Value> = browser_resp
            .get("result")
            .and_then(|r| r.get("resources"))
            .and_then(Value::as_array)
            .cloned()
            .unwrap_or_default();

        // Collect proxy resources (pool stats).
        let proxy_resp = self.proxy.handle_request(&list_req).await;
        let proxy_resources: Vec<Value> = proxy_resp
            .get("result")
            .and_then(|r| r.get("resources"))
            .and_then(Value::as_array)
            .cloned()
            .unwrap_or_default();

        let all: Vec<Value> = [browser_resources, proxy_resources]
            .into_iter()
            .flatten()
            .collect();

        ok_response(id, &json!({ "resources": all }))
    }

    // ── resources/read ────────────────────────────────────────────────────────

    async fn handle_resources_read(&self, id: &Value, req: &Value) -> Value {
        let uri = req
            .get("params")
            .and_then(|p| p.get("uri"))
            .and_then(Value::as_str)
            .unwrap_or("");

        if uri.starts_with("browser://") {
            self.browser.dispatch(req).await
        } else if uri.starts_with("proxy://") {
            self.proxy.handle_request(req).await
        } else {
            error_response(id, -32602, &format!("Unknown resource URI: {uri}"))
        }
    }

    // ── Cross-crate tool: scrape_proxied ─────────────────────────────────────

    async fn tool_scrape_proxied(&self, id: &Value, args: &Value) -> Value {
        let Some(url) = args.get("url").and_then(Value::as_str).map(str::to_owned) else {
            return error_response(id, -32602, "Missing 'url'");
        };
        let timeout_secs = args
            .get("timeout_secs")
            .and_then(Value::as_u64)
            .unwrap_or(30);

        // 1. Acquire a proxy from the pool.
        let (handle_token, proxy_url) = match self.acquire_proxy_and_token(id).await {
            Ok(v) => v,
            Err(e) => return e,
        };

        // 2. Scrape via the acquired proxy, forwarding the caller's timeout.
        let scrape_resp = McpGraphServer::handle_request(&json!({
            "jsonrpc": "2.0", "id": 0, "method": "tools/call",
            "params": {
                "name": "scrape",
                "arguments": { "url": url, "proxy_url": proxy_url, "timeout_secs": timeout_secs }
            }
        }))
        .await;
        let success = scrape_resp.get("error").is_none_or(Value::is_null);

        // 3. Release the proxy handle (mark success/failure for circuit-breaker).
        release_proxy(&self.proxy, &handle_token, success).await;

        // 4. Propagate scrape result.
        if success {
            let text = mcp_content_text_raw(&scrape_resp);
            ok_response(id, &json!({ "content": [{"type": "text", "text": text}] }))
        } else {
            let code = scrape_resp
                .get("error")
                .and_then(|e| e.get("code"))
                .and_then(Value::as_i64)
                .and_then(|c| i32::try_from(c).ok())
                .unwrap_or(-32603);
            let message = scrape_resp
                .get("error")
                .and_then(|e| e.get("message"))
                .and_then(Value::as_str)
                .unwrap_or("Graph scrape failed");
            error_response(id, code, message)
        }
    }

    // ── Cross-crate tool: browser_proxied ────────────────────────────────────

    async fn tool_browser_proxied(&self, id: &Value, args: &Value) -> Value {
        let Some(url) = args.get("url").and_then(Value::as_str).map(str::to_owned) else {
            return error_response(id, -32602, "Missing 'url'");
        };

        // 1. Acquire a proxy.
        let (handle_token, proxy_url) = match self.acquire_proxy_and_token(id).await {
            Ok(v) => v,
            Err(e) => return e,
        };

        // 2. Acquire a browser session. Note: the `proxy` argument is stored as session
        //    metadata in stygian-browser but does not yet route network traffic through the
        //    proxy at browser-launch level. Pass it for forward-compatibility.
        let acquire_browser_resp = self
            .browser
            .dispatch(&json!({
                "jsonrpc": "2.0", "id": 0, "method": "tools/call",
                "params": { "name": "browser_acquire", "arguments": { "proxy": proxy_url } }
            }))
            .await;
        // MCP tool failures from stygian-browser appear as result.isError=true.
        let acquire_is_error = acquire_browser_resp
            .get("result")
            .and_then(|r| r.get("isError"))
            .and_then(Value::as_bool)
            == Some(true);
        let session_info = parse_content_text(&acquire_browser_resp);
        let Some(session_id) = session_info
            .get("session_id")
            .and_then(Value::as_str)
            .map(str::to_owned)
        else {
            release_proxy(&self.proxy, &handle_token, false).await;
            let err_msg = if acquire_is_error {
                mcp_error_message(&acquire_browser_resp)
                    .unwrap_or("Failed to acquire browser session")
            } else {
                "Failed to acquire browser session"
            };
            return error_response(id, -32603, err_msg);
        };

        // 3. Navigate to URL.
        let nav_resp = self
            .browser
            .dispatch(&json!({
                "jsonrpc": "2.0", "id": 0, "method": "tools/call",
                "params": {
                    "name": "browser_navigate",
                    "arguments": { "session_id": session_id, "url": url }
                }
            }))
            .await;
        // Short-circuit: if navigation failed, release resources and propagate the error.
        if is_mcp_error_or_tool_error(&nav_resp) {
            self.browser_release(&session_id).await;
            release_proxy(&self.proxy, &handle_token, false).await;
            let nav_err = mcp_error_message(&nav_resp).unwrap_or("Browser navigation failed");
            return error_response(id, -32603, nav_err);
        }

        // 4. Capture HTML content.
        let content_resp = self
            .browser
            .dispatch(&json!({
                "jsonrpc": "2.0", "id": 0, "method": "tools/call",
                "params": { "name": "browser_content", "arguments": { "session_id": session_id } }
            }))
            .await;
        let content_ok = !is_mcp_error_or_tool_error(&content_resp);

        // 5. Release browser session and proxy — success only if content succeeded.
        self.browser_release(&session_id).await;
        release_proxy(&self.proxy, &handle_token, content_ok).await;

        if !content_ok {
            let content_err =
                mcp_error_message(&content_resp).unwrap_or("Browser content retrieval failed");
            return error_response(id, -32603, content_err);
        }

        // 6. Return combined result — parse sub-tool text fields as JSON to avoid
        //    double-encoded strings in the aggregated output.
        let nav_raw = mcp_content_text_raw(&nav_resp);
        let html_raw = mcp_content_text_raw(&content_resp);
        let nav_json = nav_raw
            .as_str()
            .and_then(|s| serde_json::from_str::<Value>(s).ok())
            .unwrap_or_else(|| nav_raw.clone());
        let html_json = html_raw
            .as_str()
            .and_then(|s| serde_json::from_str::<Value>(s).ok())
            .unwrap_or_else(|| html_raw.clone());

        ok_response(
            id,
            &json!({
                "content": [{
                    "type": "text",
                    "text": serde_json::to_string(&json!({
                        "navigation": nav_json,
                        "html":       html_json
                    }))
                    .unwrap_or_default()
                }]
            }),
        )
    }

    // ── Cross-crate tool: scrape_with_plugin_fallback ─────────────────────────

    /// Execute the HTTP→plugin fallback chain for a URL.
    ///
    /// Tries a plain HTTP scrape first.  If that fails or the HTTP circuit is
    /// open, the chain automatically falls back to plugin template extraction
    /// using `template_id` from `args`.
    async fn tool_scrape_with_plugin_fallback(&self, id: &Value, args: &Value) -> Value {
        let Some(url) = args.get("url").and_then(Value::as_str).map(str::to_owned) else {
            return error_response(id, -32602, "Missing 'url'");
        };
        let Some(template_id) = args
            .get("template_id")
            .and_then(Value::as_str)
            .map(str::to_owned)
        else {
            return error_response(id, -32602, "Missing 'template_id'");
        };

        let mut params = json!({ "template_id": template_id });

        // Include HTML if provided, so plugin fallback can use it directly
        if let Some(html) = args.get("html").and_then(Value::as_str)
            && let Some(obj) = params.as_object_mut()
        {
            obj.insert("html".to_string(), Value::String(html.to_owned()));
        }

        if let Some(idem_key) = args.get("idempotency_key").and_then(Value::as_str)
            && let Some(obj) = params.as_object_mut()
        {
            obj.insert(
                "idempotency_key".to_string(),
                Value::String(idem_key.to_owned()),
            );
        }

        let input = ServiceInput { url, params };
        match self.fallback_chain.execute(input).await {
            Ok(output) => ok_response(
                id,
                &json!({
                    "content": [{
                        "type": "text",
                        "text": output.data
                    }],
                    "metadata": output.metadata
                }),
            ),
            Err(e) => error_response(id, -32603, &format!("Fallback chain exhausted: {e}")),
        }
    }

    // ── Private helpers ───────────────────────────────────────────────────────

    /// Acquire a proxy from the pool, returning `(handle_token, proxy_url)` on success
    /// or a ready-made JSON-RPC error response on the `Err` side.
    async fn acquire_proxy_and_token(&self, id: &Value) -> Result<(String, String), Value> {
        let acquire_resp = self
            .proxy
            .handle_request(&json!({
                "jsonrpc": "2.0", "id": 0, "method": "tools/call",
                "params": { "name": "proxy_acquire", "arguments": {} }
            }))
            .await;
        if let Some(err) = acquire_resp.get("error").filter(|e| !e.is_null()) {
            let code = err
                .get("code")
                .and_then(Value::as_i64)
                .and_then(|c| i32::try_from(c).ok())
                .unwrap_or(-32603);
            let message = err
                .get("message")
                .and_then(Value::as_str)
                .unwrap_or("No proxy available — add proxies via proxy_add first");
            return Err(error_response(id, code, message));
        }
        let handle_info = parse_content_text(&acquire_resp);
        let Some(handle_token) = handle_info
            .get("handle_token")
            .and_then(Value::as_str)
            .map(str::to_owned)
        else {
            return Err(error_response(
                id,
                -32603,
                "No proxy available — add proxies via proxy_add first",
            ));
        };
        let Some(proxy_url) = handle_info
            .get("proxy_url")
            .and_then(Value::as_str)
            .map(str::to_owned)
        else {
            release_proxy(&self.proxy, &handle_token, false).await;
            return Err(error_response(
                id,
                -32603,
                "proxy_acquire returned no proxy_url",
            ));
        };
        Ok((handle_token, proxy_url))
    }

    /// Release a browser session via the browser sub-server.
    async fn browser_release(&self, session_id: &str) {
        let _ = self
            .browser
            .dispatch(&json!({
                "jsonrpc": "2.0", "id": 0, "method": "tools/call",
                "params": {
                    "name": "browser_release",
                    "arguments": { "session_id": session_id }
                }
            }))
            .await;
    }
}

impl Drop for McpAggregator {
    /// Cancel and abort the proxy background tasks on any drop path so they
    /// do not outlive the aggregator when `run()` returns early via `?`.
    fn drop(&mut self) {
        self.proxy_token.cancel();
        if let Some(bg) = self.proxy_bg.take() {
            bg.abort();
        }
    }
}

// ─── Helpers ─────────────────────────────────────────────────────────────────

fn handle_initialize(req: &Value) -> Value {
    let id = req.get("id").unwrap_or(&Value::Null);
    let requested = req
        .get("params")
        .and_then(|p| p.get("protocolVersion"))
        .and_then(Value::as_str);

    let protocol_version = match requested {
        Some(version) if SUPPORTED_PROTOCOL_VERSIONS.contains(&version) => version,
        Some(version) => {
            return error_response(
                id,
                -32602,
                &format!(
                    "Unsupported protocolVersion: {version}. Supported: {}",
                    SUPPORTED_PROTOCOL_VERSIONS.join(", ")
                ),
            );
        }
        None => SUPPORTED_PROTOCOL_VERSIONS
            .first()
            .copied()
            .unwrap_or("2024-11-05"),
    };

    ok_response(
        id,
        &json!({
            "protocolVersion": protocol_version,
            "capabilities": {
                "tools":     { "listChanged": false },
                "resources": { "listChanged": false, "subscribe": false }
            },
            "serverInfo": {
                "name":    "stygian-mcp",
                "version": env!("CARGO_PKG_VERSION")
            }
        }),
    )
}

fn ok_response(id: &Value, result: &Value) -> Value {
    json!({ "jsonrpc": "2.0", "id": id, "result": result })
}

fn error_response(id: &Value, code: i32, message: &str) -> Value {
    json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": { "code": code, "message": message }
    })
}

/// Extract the first content text item from an MCP response and parse it as
/// JSON.  Falls back to `Value::Null` on any parsing failure.
fn parse_content_text(resp: &Value) -> Value {
    resp.get("result")
        .and_then(|r| r.get("content"))
        .and_then(|c| c.get(0))
        .and_then(|c| c.get("text"))
        .and_then(Value::as_str)
        .and_then(|s| serde_json::from_str(s).ok())
        .unwrap_or(Value::Null)
}

/// Extract the raw first-content-item text from an MCP response as an owned `Value`.
///
/// Returns `Value::Null` if no content text is present.
fn mcp_content_text_raw(resp: &Value) -> Value {
    resp.get("result")
        .and_then(|r| r.get("content"))
        .and_then(|c| c.get(0))
        .and_then(|c| c.get("text"))
        .cloned()
        .unwrap_or(Value::Null)
}

/// Returns `true` if the response has a non-null `error` field or `result.isError = true`.
fn is_mcp_error_or_tool_error(resp: &Value) -> bool {
    !resp.get("error").is_none_or(Value::is_null)
        || resp
            .get("result")
            .and_then(|r| r.get("isError"))
            .and_then(Value::as_bool)
            == Some(true)
}

/// Extract the error message from a JSON-RPC / MCP tool-error response.
///
/// Checks `error.message` first; falls back to `result.content[0].text`
/// (the `isError` path used by `stygian-browser`).
fn mcp_error_message(resp: &Value) -> Option<&str> {
    resp.get("error")
        .filter(|e| !e.is_null())
        .and_then(|e| e.get("message"))
        .and_then(Value::as_str)
        .or_else(|| {
            resp.get("result")
                .and_then(|r| r.get("content"))
                .and_then(|c| c.get(0))
                .and_then(|c| c.get("text"))
                .and_then(Value::as_str)
        })
}

fn is_jsonrpc_notification(req: &Value) -> bool {
    req.is_object()
        && req.get("jsonrpc").and_then(Value::as_str) == Some("2.0")
        && req.get("id").is_none()
        && req.get("method").and_then(Value::as_str).is_some()
}

/// Enrich `browser_verify_stealth` responses with an aggregation summary while
/// preserving the underlying `DiagnosticReport` payload.
fn normalize_browser_verify_stealth_response(mut resp: Value) -> Value {
    let Some(text_slot) = resp
        .get_mut("result")
        .and_then(Value::as_object_mut)
        .and_then(|r| r.get_mut("content"))
        .and_then(Value::as_array_mut)
        .and_then(|a| a.get_mut(0))
        .and_then(Value::as_object_mut)
        .and_then(|c| c.get_mut("text"))
    else {
        return resp;
    };

    let Some(raw_text) = text_slot.as_str() else {
        return resp;
    };

    let Ok(mut report) = serde_json::from_str::<Value>(raw_text) else {
        return resp;
    };

    let failed_check_ids: Vec<Value> = report
        .get("checks")
        .and_then(Value::as_array)
        .into_iter()
        .flatten()
        .filter(|check| check.get("passed").and_then(Value::as_bool) == Some(false))
        .map(|check| {
            check
                .get("id")
                .cloned()
                .unwrap_or_else(|| Value::String("unknown".to_string()))
        })
        .collect();

    let known_limitation_ids: Vec<Value> = report
        .get("known_limitations")
        .and_then(Value::as_array)
        .into_iter()
        .flatten()
        .map(|limitation| {
            limitation
                .get("id")
                .cloned()
                .unwrap_or_else(|| Value::String("unknown".to_string()))
        })
        .collect();

    let known_limitations_count = u64::try_from(known_limitation_ids.len()).unwrap_or(0);
    let failed_checks_count = u64::try_from(failed_check_ids.len()).unwrap_or(0);

    if let Some(report_obj) = report.as_object_mut() {
        report_obj.insert(
            "aggregation_summary".to_string(),
            json!({
                "failed_checks_count": failed_checks_count,
                "failed_check_ids": failed_check_ids,
                "known_limitations_count": known_limitations_count,
                "known_limitation_ids": known_limitation_ids
            }),
        );
    }

    if let Ok(serialized) = serde_json::to_string(&report) {
        *text_slot = Value::String(serialized);
    }

    resp
}

/// Release a proxy handle via the proxy sub-server.
async fn release_proxy(proxy: &Arc<McpProxyServer>, handle_token: &str, success: bool) {
    let _ = proxy
        .handle_request(&json!({
            "jsonrpc": "2.0", "id": 0, "method": "tools/call",
            "params": {
                "name": "proxy_release",
                "arguments": { "handle_token": handle_token, "success": success }
            }
        }))
        .await;
}

// ─── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use stygian_graph::mcp::McpGraphServer;
    use stygian_proxy::mcp::McpProxyServer;

    /// Validate that the aggregator routes `graph_scrape` to the graph server.
    #[tokio::test]
    async fn test_routes_graph_tool() -> Result<(), Box<dyn std::error::Error>> {
        let proxy = Arc::new(McpProxyServer::new()?);
        let list_req = json!({
            "jsonrpc": "2.0", "id": 1,
            "method": "tools/list",
            "params": {}
        });
        let resp = McpGraphServer::handle_request(&list_req).await;
        let tools = resp
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(Value::as_array)
            .ok_or("graph tools/list missing result.tools")?;
        // Confirm graph server exposes `scrape` (un-prefixed).
        assert!(
            tools
                .iter()
                .any(|t| t.get("name").and_then(Value::as_str) == Some("scrape"))
        );
        // Proxy list should not contain `scrape`.
        let proxy_resp = proxy.handle_request(&list_req).await;
        let proxy_tools = proxy_resp
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(Value::as_array)
            .ok_or("proxy tools/list missing result.tools")?;
        assert!(
            !proxy_tools
                .iter()
                .any(|t| t.get("name").and_then(Value::as_str) == Some("scrape"))
        );
        drop(proxy);
        Ok(())
    }

    #[test]
    fn test_error_response_structure() {
        let resp = error_response(&json!(42), -32602, "bad param");
        assert_eq!(
            resp.pointer("/error/code").and_then(Value::as_i64),
            Some(-32602)
        );
        assert_eq!(resp.get("id").and_then(Value::as_i64), Some(42));
    }

    #[test]
    fn test_ok_response_structure() {
        let resp = ok_response(&json!(1), &json!({"foo": "bar"}));
        assert_eq!(
            resp.pointer("/result/foo").and_then(Value::as_str),
            Some("bar")
        );
        assert_eq!(resp.get("jsonrpc").and_then(Value::as_str), Some("2.0"));
    }

    #[test]
    fn test_normalize_browser_verify_stealth_response_adds_summary() {
        let resp = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "result": {
                "content": [{
                    "type": "text",
                    "text": "{\"checks\":[{\"id\":\"web_driver_flag\",\"passed\":false},{\"id\":\"chrome_object\",\"passed\":true}],\"known_limitations\":[{\"id\":\"web_gpu_surface\"}],\"passed_count\":1,\"failed_count\":1}"
                }]
            }
        });

        let normalized = normalize_browser_verify_stealth_response(resp);
        let text = normalized
            .pointer("/result/content/0/text")
            .and_then(Value::as_str)
            .unwrap_or("");
        let parsed: Value = serde_json::from_str(text).unwrap_or(Value::Null);

        assert_eq!(
            parsed
                .pointer("/aggregation_summary/failed_checks_count")
                .and_then(Value::as_u64),
            Some(1)
        );
        assert_eq!(
            parsed
                .pointer("/aggregation_summary/known_limitations_count")
                .and_then(Value::as_u64),
            Some(1)
        );
        assert_eq!(
            parsed
                .pointer("/aggregation_summary/failed_check_ids/0")
                .and_then(Value::as_str),
            Some("web_driver_flag")
        );
        assert_eq!(
            parsed
                .pointer("/aggregation_summary/known_limitation_ids/0")
                .and_then(Value::as_str),
            Some("web_gpu_surface")
        );
    }

    #[test]
    fn test_normalize_browser_verify_stealth_response_passthrough_on_non_json_text() {
        let resp = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "result": {
                "content": [{
                    "type": "text",
                    "text": "not-json"
                }]
            }
        });

        let normalized = normalize_browser_verify_stealth_response(resp.clone());
        assert_eq!(normalized, resp);
    }

    #[test]
    fn test_initialize_negotiates_supported_protocol() {
        let req = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": { "protocolVersion": "2024-11-05" }
        });
        let resp = handle_initialize(&req);
        assert_eq!(
            resp.pointer("/result/protocolVersion")
                .and_then(Value::as_str),
            Some("2024-11-05")
        );
    }

    #[test]
    fn test_initialize_rejects_unsupported_protocol() {
        let req = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": { "protocolVersion": "1999-01-01" }
        });
        let resp = handle_initialize(&req);
        assert_eq!(
            resp.pointer("/error/code").and_then(Value::as_i64),
            Some(-32602)
        );
    }

    /// `scrape_proxied` with no proxies in the pool must return an error
    /// (not silently succeed or panic).
    #[tokio::test]
    async fn test_scrape_proxied_no_proxy_returns_error() -> Result<(), Box<dyn std::error::Error>>
    {
        let proxy = Arc::new(McpProxyServer::new()?);
        let id = json!(99);
        let acquire_req = json!({
            "jsonrpc": "2.0", "id": 0, "method": "tools/call",
            "params": { "name": "proxy_acquire", "arguments": {} }
        });
        let acquire_resp = proxy.handle_request(&acquire_req).await;
        // Simulate the scrape_proxied acquire-failure branch.
        let resp = if acquire_resp.get("error").is_none_or(Value::is_null) {
            // parse_content_text returns Null for empty pool → no handle_token
            let handle_info = parse_content_text(&acquire_resp);
            if handle_info
                .get("handle_token")
                .and_then(Value::as_str)
                .is_none()
            {
                error_response(
                    &id,
                    -32603,
                    "No proxy available — add proxies via proxy_add first",
                )
            } else {
                ok_response(&id, &json!({"unexpected": "success"}))
            }
        } else {
            let code = acquire_resp
                .get("error")
                .and_then(|e| e.get("code"))
                .and_then(Value::as_i64)
                .and_then(|c| i32::try_from(c).ok())
                .unwrap_or(-32603);
            let message = acquire_resp
                .get("error")
                .and_then(|e| e.get("message"))
                .and_then(Value::as_str)
                .unwrap_or("No proxy available");
            error_response(&id, code, message)
        };
        assert!(
            resp.get("error").is_some_and(|e| !e.is_null()),
            "expected an error response"
        );
        assert_eq!(
            resp.get("id").and_then(Value::as_i64),
            Some(99),
            "id must be the caller's id, not 0"
        );
        drop(proxy);
        Ok(())
    }

    /// `browser_proxied` must propagate a JSON-RPC error with the correct caller id
    /// when proxy acquisition fails (empty pool).
    #[tokio::test]
    async fn test_browser_proxied_acquire_failure_uses_caller_id()
    -> Result<(), Box<dyn std::error::Error>> {
        let proxy = Arc::new(McpProxyServer::new()?);
        let id = json!(77);
        let acquire_req = json!({
            "jsonrpc": "2.0", "id": 0, "method": "tools/call",
            "params": { "name": "proxy_acquire", "arguments": {} }
        });
        let acquire_resp = proxy.handle_request(&acquire_req).await;
        // With an empty pool the acquire either errors or returns no handle_token.
        // Either way the caller's id must survive in the response.
        let resp = if acquire_resp.get("error").is_none_or(Value::is_null) {
            let handle_info = parse_content_text(&acquire_resp);
            if handle_info
                .get("handle_token")
                .and_then(Value::as_str)
                .is_none()
            {
                error_response(
                    &id,
                    -32603,
                    "No proxy available — add proxies via proxy_add first",
                )
            } else {
                ok_response(&id, &json!({"unexpected": "success"}))
            }
        } else {
            let code = acquire_resp
                .get("error")
                .and_then(|e| e.get("code"))
                .and_then(Value::as_i64)
                .and_then(|c| i32::try_from(c).ok())
                .unwrap_or(-32603);
            let message = acquire_resp
                .get("error")
                .and_then(|e| e.get("message"))
                .and_then(Value::as_str)
                .unwrap_or("No proxy available");
            error_response(&id, code, message)
        };
        assert!(
            resp.get("error").is_some_and(|e| !e.is_null()),
            "expected an error response"
        );
        // The critical invariant: id must be 77, not the internal sub-request id 0.
        assert_eq!(
            resp.get("id").and_then(Value::as_i64),
            Some(77),
            "caller id must be preserved in error response"
        );
        Ok(())
    }

    /// Verify that `parse_content_text` returns Null on a JSON-RPC error response
    /// (no `result.content` field), which is the sentinel for "acquire failed".
    #[test]
    fn test_parse_content_text_on_error_response_returns_null() {
        let error_resp = error_response(&json!(1), -32603, "internal error");
        let parsed = parse_content_text(&error_resp);
        assert_eq!(parsed, Value::Null);
    }

    #[test]
    fn test_notification_detection_requires_valid_jsonrpc() {
        let valid_notification = json!({
            "jsonrpc": "2.0",
            "method": "notifications/initialized",
            "params": {}
        });
        assert!(is_jsonrpc_notification(&valid_notification));

        let missing_jsonrpc = json!({
            "method": "notifications/initialized",
            "params": {}
        });
        assert!(!is_jsonrpc_notification(&missing_jsonrpc));
    }

    #[test]
    fn test_invalid_request_missing_jsonrpc_returns_invalid_request_error() {
        let req = json!({
            "method": "tools/list"
        });
        let id = req.get("id").unwrap_or(&Value::Null);
        let resp = if req.get("jsonrpc").and_then(Value::as_str) == Some("2.0") {
            ok_response(id, &json!({}))
        } else {
            error_response(id, -32600, "Invalid request: expected jsonrpc='2.0'")
        };
        assert_eq!(
            resp.pointer("/error/code").and_then(Value::as_i64),
            Some(-32600)
        );
        assert!(resp.get("id").is_none_or(Value::is_null));
    }

    /// `tools/list` aggregation: graph tools get a `graph_` prefix and a `[graph]` description
    /// prefix; proxy tools keep their existing `proxy_` prefix; all lists are merged.
    #[tokio::test]
    async fn test_tools_list_prefixes_graph_tools() -> Result<(), Box<dyn std::error::Error>> {
        let proxy = Arc::new(McpProxyServer::new()?);

        let list_req = json!({ "jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {} });

        // Graph: un-prefixed names like `scrape`
        let graph_resp = McpGraphServer::handle_request(&list_req).await;
        let graph_tools = graph_resp
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(Value::as_array)
            .ok_or("graph tools/list missing result.tools")?;
        assert!(
            graph_tools
                .iter()
                .any(|t| t.get("name").and_then(Value::as_str) == Some("scrape")),
            "graph server exposes un-prefixed 'scrape'"
        );

        // Proxy: already prefixed
        let proxy_resp = proxy.handle_request(&list_req).await;
        let proxy_tools = proxy_resp
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(Value::as_array)
            .ok_or("proxy tools/list missing result.tools")?;
        assert!(
            proxy_tools
                .iter()
                .any(|t| t.get("name").and_then(Value::as_str) == Some("proxy_add")),
            "proxy server exposes 'proxy_add'"
        );

        // Simulate the prefixing logic from handle_tools_list
        let prefixed: Vec<String> = graph_tools
            .iter()
            .filter_map(|t| t.get("name").and_then(Value::as_str))
            .map(|n| format!("graph_{n}"))
            .collect();
        assert!(
            prefixed.contains(&"graph_scrape".to_string()),
            "graph_scrape must appear after prefixing"
        );
        assert!(
            !prefixed.contains(&"scrape".to_string()),
            "un-prefixed 'scrape' must not appear after prefixing"
        );
        Ok(())
    }

    /// `tools/call` dispatch: names with `graph_` prefix are routed to the graph server with
    /// the prefix stripped; `proxy_` names are routed to the proxy server unchanged.
    #[tokio::test]
    async fn test_tools_call_dispatch_by_prefix() -> Result<(), Box<dyn std::error::Error>> {
        let proxy = Arc::new(McpProxyServer::new()?);

        // graph_pipeline_validate → graph server with name `pipeline_validate`
        let graph_call = json!({
            "jsonrpc": "2.0", "id": 2, "method": "tools/call",
            "params": { "name": "pipeline_validate", "arguments": { "toml": "" } }
        });
        let graph_resp = McpGraphServer::handle_request(&graph_call).await;
        // pipeline_validate with empty string returns a result (not an unknown-method error)
        assert!(
            graph_resp.get("error").is_none_or(Value::is_null)
                || graph_resp
                    .pointer("/result/content/0/text")
                    .and_then(Value::as_str)
                    .is_some(),
            "graph server must respond to pipeline_validate"
        );

        // proxy_add with missing url returns an error from the proxy server (not unknown tool)
        let proxy_call = json!({
            "jsonrpc": "2.0", "id": 3, "method": "tools/call",
            "params": { "name": "proxy_add", "arguments": {} }
        });
        let proxy_resp = proxy.handle_request(&proxy_call).await;
        // Should get an error about missing `url`, not about an unknown tool
        let err_msg = proxy_resp
            .get("error")
            .and_then(|e| e.get("message"))
            .and_then(Value::as_str)
            .or_else(|| {
                proxy_resp
                    .pointer("/result/content/0/text")
                    .and_then(Value::as_str)
            })
            .unwrap_or("");
        assert!(
            err_msg.contains("url") || err_msg.contains("required") || !err_msg.is_empty(),
            "proxy server must respond to proxy_add (got: {err_msg})"
        );
        Ok(())
    }

    /// `resources/list` aggregation: proxy resources (pool stats) must be collected.
    #[tokio::test]
    async fn test_resources_list_includes_proxy() -> Result<(), Box<dyn std::error::Error>> {
        let proxy = Arc::new(McpProxyServer::new()?);

        let list_req =
            json!({ "jsonrpc": "2.0", "id": 4, "method": "resources/list", "params": {} });
        let proxy_resp = proxy.handle_request(&list_req).await;
        let resources = proxy_resp
            .get("result")
            .and_then(|r| r.get("resources"))
            .and_then(Value::as_array);
        // The proxy pool stats resource is always present at proxy://pool/stats
        assert!(
            resources.is_some_and(|r| r
                .iter()
                .any(|res| res.get("uri").and_then(Value::as_str) == Some("proxy://pool/stats"))),
            "proxy resources/list must contain proxy://pool/stats"
        );
        Ok(())
    }
}
