//! MCP JSON-RPC 2.0 server backed by `provisioning-core::Registry`.
//!
//! Transport: newline-delimited JSON on stdin/stdout (same as the legacy
//! `simple_main.rs`). Each line is one JSON-RPC 2.0 message.
//!
//! Supported methods:
//!   initialize          — MCP handshake → server capabilities
//!   tools/list          — enumerate all registered tools
//!   tools/call          — invoke a tool by name
//!   ping                — no-op health probe
//!
//! Notifications (no `id` field) are silently dropped.

use provisioning_core::{Environment, Registry, ToolError, tool::Context};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::io::{BufRead, Write};
use std::sync::Arc;

const PROTOCOL_VERSION: &str = "2024-11-05";

// ── JSON-RPC types ────────────────────────────────────────────────────────────

#[derive(Debug, Deserialize)]
struct RpcRequest {
    #[allow(dead_code)]
    jsonrpc: Option<String>,
    id: Option<Value>,
    method: String,
    params: Option<Value>,
}

#[derive(Debug, Serialize)]
struct RpcResponse {
    jsonrpc: &'static str,
    id: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<RpcError>,
}

#[derive(Debug, Serialize)]
struct RpcError {
    code: i32,
    message: String,
}

impl RpcResponse {
    fn ok(id: Value, result: Value) -> Self {
        Self { jsonrpc: "2.0", id, result: Some(result), error: None }
    }

    fn err(id: Value, code: i32, message: impl Into<String>) -> Self {
        Self {
            jsonrpc: "2.0",
            id,
            result: None,
            error: Some(RpcError { code, message: message.into() }),
        }
    }
}

// ── McpServer ─────────────────────────────────────────────────────────────────

pub struct McpServer {
    registry: Arc<Registry>,
    env: Arc<Environment>,
    pub server_name: String,
    pub server_version: String,
}

impl McpServer {
    pub fn new(
        registry: Registry,
        env: Environment,
        server_name: impl Into<String>,
        server_version: impl Into<String>,
    ) -> Self {
        Self {
            registry: Arc::new(registry),
            env: Arc::new(env),
            server_name: server_name.into(),
            server_version: server_version.into(),
        }
    }

    /// Run the server loop on stdin/stdout until EOF.
    pub async fn run_stdio(&self) -> anyhow::Result<()> {
        let stdin = std::io::stdin();
        let mut stdout = std::io::stdout();

        for line in stdin.lock().lines() {
            let line = line?;
            if line.trim().is_empty() {
                continue;
            }

            let req: RpcRequest = match serde_json::from_str(&line) {
                Ok(r) => r,
                Err(e) => {
                    let resp = RpcResponse::err(Value::Null, -32700, format!("parse error: {e}"));
                    writeln!(stdout, "{}", serde_json::to_string(&resp)?)?;
                    stdout.flush()?;
                    continue;
                }
            };

            // Notifications (no id) → silently drop
            if req.id.is_none() {
                continue;
            }

            let id = req.id.unwrap();
            let resp = self.dispatch(id, &req.method, req.params).await;
            writeln!(stdout, "{}", serde_json::to_string(&resp)?)?;
            stdout.flush()?;
        }

        Ok(())
    }

    /// In-process entry point for contract tests — accepts a raw JSON-RPC
    /// envelope, returns the response as a `Value`. No stdio involvement.
    ///
    /// Used by the G3 contract test to invoke the MCP surface without spawning
    /// a subprocess or touching stdin/stdout.
    pub async fn handle_request(&self, request: Value) -> Value {
        let req: RpcRequest = match serde_json::from_value(request.clone()) {
            Ok(r) => r,
            Err(e) => {
                let resp = RpcResponse::err(Value::Null, -32700, format!("parse error: {e}"));
                return serde_json::to_value(resp).unwrap_or(Value::Null);
            }
        };
        let id = req.id.unwrap_or(Value::Null);
        let resp = self.dispatch(id, &req.method, req.params).await;
        serde_json::to_value(resp).unwrap_or(Value::Null)
    }

    async fn dispatch(&self, id: Value, method: &str, params: Option<Value>) -> RpcResponse {
        match method {
            "initialize" => self.handle_initialize(id),
            "tools/list" => self.handle_tools_list(id),
            "tools/call" => self.handle_tools_call(id, params).await,
            "ping" => RpcResponse::ok(id, json!({})),
            other => RpcResponse::err(id, -32601, format!("method not found: {other}")),
        }
    }

    fn handle_initialize(&self, id: Value) -> RpcResponse {
        RpcResponse::ok(id, json!({
            "protocolVersion": PROTOCOL_VERSION,
            "capabilities": {
                "tools": { "listChanged": false }
            },
            "serverInfo": {
                "name":    self.server_name,
                "version": self.server_version,
            }
        }))
    }

    fn handle_tools_list(&self, id: Value) -> RpcResponse {
        let tools: Vec<Value> = self
            .registry
            .list()
            .into_iter()
            .map(|m| json!({
                "name":        m.name,
                "description": m.description,
                "inputSchema": *m.schema,
            }))
            .collect();

        RpcResponse::ok(id, json!({ "tools": tools }))
    }

    async fn handle_tools_call(&self, id: Value, params: Option<Value>) -> RpcResponse {
        let params = params.unwrap_or(Value::Null);
        let name = match params["name"].as_str() {
            Some(n) => n,
            None => return RpcResponse::err(id, -32602, "params.name is required"),
        };
        let args = params["arguments"].clone();

        let ctx = Context::new(Arc::clone(&self.env));
        match self.registry.invoke(name, args, &ctx).await {
            Ok(result) => {
                let text = serde_json::to_string_pretty(&result)
                    .unwrap_or_else(|_| result.to_string());
                RpcResponse::ok(id, json!({
                    "content": [{ "type": "text", "text": text }]
                }))
            }
            Err(e) => {
                let (code, msg) = tool_error_to_rpc(&e);
                RpcResponse::err(id, code, msg)
            }
        }
    }
}

// ── Error mapping ─────────────────────────────────────────────────────────────

fn tool_error_to_rpc(e: &ToolError) -> (i32, String) {
    match e {
        ToolError::NotFound(_) => (-32001, e.to_string()),
        ToolError::InvalidParam { .. } => (-32602, e.to_string()),
        ToolError::Unauthorized(_) => (-32003, e.to_string()),
        _ => (-32000, e.to_string()),
    }
}

// ── Tool catalog helper (used by D2 contract test) ────────────────────────────

/// Generate the `tools/list` result payload without running the server loop.
pub fn tools_list_payload(registry: &Registry) -> Value {
    let tools: Vec<Value> = registry
        .list()
        .into_iter()
        .map(|m| json!({
            "name":        m.name,
            "description": m.description,
            "inputSchema": *m.schema,
        }))
        .collect();
    json!({ "tools": tools })
}

#[cfg(test)]
mod tests {
    use super::*;
    use provisioning_core::{
        sources::{NclCache, OrchestratorClient},
        providers::hcloud::HcloudClient,
    };

    fn make_registry() -> Registry {
        let cache  = Arc::new(NclCache::new());
        let orch   = Arc::new(OrchestratorClient::new("http://localhost:19999"));
        let hcloud = Arc::new(HcloudClient::new());
        Registry::with_all_tools(cache, orch, hcloud, None).unwrap()
    }

    #[test]
    fn tools_list_payload_contains_all_registry_tools() {
        let reg = make_registry();
        let expected_count = reg.len();
        let payload = tools_list_payload(&reg);

        let tools = payload["tools"].as_array().expect("tools must be an array");
        assert_eq!(
            tools.len(),
            expected_count,
            "MCP tools/list count ({}) != registry count ({})",
            tools.len(),
            expected_count
        );
    }

    #[test]
    fn tools_list_each_tool_has_name_description_schema() {
        let reg = make_registry();
        let payload = tools_list_payload(&reg);
        let tools = payload["tools"].as_array().unwrap();

        for tool in tools {
            assert!(tool["name"].is_string(), "tool missing name: {tool}");
            assert!(tool["description"].is_string(), "tool missing description: {tool}");
            assert_eq!(tool["inputSchema"]["type"], "object", "tool schema must be object: {}", tool["name"]);
        }
    }

    #[test]
    fn tools_list_names_match_registry_names() {
        let reg = make_registry();
        let payload = tools_list_payload(&reg);
        let tools = payload["tools"].as_array().unwrap();

        let registry_names: std::collections::HashSet<&str> =
            reg.list().iter().map(|m| m.name).collect();
        let mcp_names: std::collections::HashSet<&str> = tools
            .iter()
            .filter_map(|t| t["name"].as_str())
            .collect();

        assert_eq!(registry_names, mcp_names, "MCP names differ from registry names");
    }

    #[test]
    fn handle_initialize_returns_protocol_version() {
        let reg = make_registry();
        let server = McpServer::new(
            reg,
            Environment::default(),
            "test-server",
            "1.0.0",
        );
        let resp = server.handle_initialize(json!(1));
        assert_eq!(resp.result.unwrap()["protocolVersion"], PROTOCOL_VERSION);
    }

    #[tokio::test]
    async fn handle_tools_call_unknown_tool_returns_not_found_error() {
        let reg = make_registry();
        let server = McpServer::new(reg, Environment::default(), "test", "0.0.1");
        let resp = server
            .dispatch(json!(1), "tools/call", Some(json!({"name": "no_such_tool", "arguments": {}})))
            .await;
        let err = resp.error.unwrap();
        assert_eq!(err.code, -32001);
    }

    #[tokio::test]
    async fn handle_tools_call_valid_tool_returns_content() {
        let reg = make_registry();
        let server = McpServer::new(reg, Environment::default(), "test", "0.0.1");
        let resp = server
            .dispatch(
                json!(1),
                "tools/call",
                Some(json!({"name": "installer_settings_defaults", "arguments": {"mode": "solo"}})),
            )
            .await;
        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
        let content = resp.result.unwrap();
        assert!(content["content"].is_array());
        assert_eq!(content["content"][0]["type"], "text");
    }

    #[tokio::test]
    async fn ping_returns_empty_object() {
        let reg = make_registry();
        let server = McpServer::new(reg, Environment::default(), "test", "0.0.1");
        let resp = server.dispatch(json!(99), "ping", None).await;
        assert!(resp.error.is_none());
        assert_eq!(resp.result.unwrap(), json!({}));
    }

    #[tokio::test]
    async fn unknown_method_returns_method_not_found() {
        let reg = make_registry();
        let server = McpServer::new(reg, Environment::default(), "test", "0.0.1");
        let resp = server.dispatch(json!(1), "nonexistent/method", None).await;
        assert_eq!(resp.error.unwrap().code, -32601);
    }
}
