//! Minimal MCP (Model Context Protocol) client over stdio — line-delimited
//! JSON-RPC 2.0. Servers are configured in ~/.lucid-train/config.toml:
//!
//! ```toml
//! [mcp_servers.github]
//! command = "npx"
//! args = ["-y", "@modelcontextprotocol/server-github"]
//! ```
//!
//! Each server tool is exposed to the model as `mcp__<server>__<tool>`.

use anyhow::{anyhow, bail, Context, Result};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, Command};
use tokio::sync::oneshot;

pub const MCP_PROTOCOL_VERSION: &str = "2024-11-05";
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);

#[derive(Debug, Clone, Default)]
pub struct McpServerConfig {
    pub name: String,
    pub command: String,
    pub args: Vec<String>,
    /// Extra environment variables passed to the server process. Required by
    /// servers like `mcp-grafana` (GRAFANA_URL, GRAFANA_SERVICE_ACCOUNT_TOKEN…).
    pub env: Vec<(String, String)>,
}

#[derive(Debug, Clone)]
pub struct McpTool {
    pub server: String,
    pub name: String,
    pub description: String,
    pub input_schema: Value,
}

impl McpTool {
    /// Name as exposed to the LLM.
    pub fn llm_name(&self) -> String {
        format!("mcp__{}__{}", self.server, self.name)
    }
}

type Pending = Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>;

struct McpConnection {
    _child: Child,
    stdin: ChildStdin,
    pending: Pending,
    next_id: AtomicU64,
}

impl McpConnection {
    async fn request(&mut self, method: &str, params: Value) -> Result<Value> {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
        let (tx, rx) = oneshot::channel();
        self.pending.lock().unwrap().insert(id, tx);
        let msg = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
        self.stdin
            .write_all((msg.to_string() + "\n").as_bytes())
            .await
            .context("mcp write failed")?;
        self.stdin.flush().await.ok();
        let resp = tokio::time::timeout(REQUEST_TIMEOUT, rx)
            .await
            .map_err(|_| anyhow!("mcp request timed out: {method}"))?
            .map_err(|_| anyhow!("mcp connection closed"))?;
        if let Some(err) = resp.get("error").filter(|e| !e.is_null()) {
            bail!("mcp error: {err}");
        }
        Ok(resp["result"].clone())
    }

    async fn notify(&mut self, method: &str, params: Value) -> Result<()> {
        let msg = json!({ "jsonrpc": "2.0", "method": method, "params": params });
        self.stdin
            .write_all((msg.to_string() + "\n").as_bytes())
            .await
            .context("mcp write failed")?;
        self.stdin.flush().await.ok();
        Ok(())
    }
}

pub struct McpManager {
    connections: HashMap<String, McpConnection>,
    pub tools: Vec<McpTool>,
    pub errors: Vec<String>,
}

impl McpManager {
    /// Spawn and initialize every configured server; tool failures are
    /// collected as warnings rather than failing the session.
    pub async fn start(configs: &[McpServerConfig]) -> McpManager {
        let mut mgr = McpManager { connections: HashMap::new(), tools: Vec::new(), errors: Vec::new() };
        for cfg in configs {
            match Self::connect(cfg).await {
                Ok((conn, tools)) => {
                    mgr.connections.insert(cfg.name.clone(), conn);
                    mgr.tools.extend(tools);
                }
                Err(e) => mgr.errors.push(format!("mcp server '{}' failed: {e:#}", cfg.name)),
            }
        }
        mgr
    }

    async fn connect(cfg: &McpServerConfig) -> Result<(McpConnection, Vec<McpTool>)> {
        let mut child = Command::new(&cfg.command)
            .args(&cfg.args)
            .envs(cfg.env.iter().map(|(k, v)| (k.as_str(), v.as_str())))
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .kill_on_drop(true)
            .spawn()
            .with_context(|| format!("failed to spawn {}", cfg.command))?;

        let stdin = child.stdin.take().context("no stdin")?;
        let stdout = child.stdout.take().context("no stdout")?;
        let pending: Pending = Arc::new(Mutex::new(HashMap::new()));

        // Reader task: route responses to pending requests by id.
        let pending_reader = pending.clone();
        tokio::spawn(async move {
            let mut lines = BufReader::new(stdout).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                let Ok(v) = serde_json::from_str::<Value>(&line) else { continue };
                if let Some(id) = v["id"].as_u64() {
                    if let Some(tx) = pending_reader.lock().unwrap().remove(&id) {
                        let _ = tx.send(v);
                    }
                }
                // Server-initiated requests/notifications are ignored (MVP).
            }
        });

        let mut conn = McpConnection { _child: child, stdin, pending, next_id: AtomicU64::new(1) };

        conn.request(
            "initialize",
            json!({
                "protocolVersion": MCP_PROTOCOL_VERSION,
                "capabilities": {},
                "clientInfo": { "name": "lucid-train", "version": env!("CARGO_PKG_VERSION") }
            }),
        )
        .await?;
        conn.notify("notifications/initialized", json!({})).await?;

        let result = conn.request("tools/list", json!({})).await?;
        let tools = result["tools"]
            .as_array()
            .map(|arr| {
                arr.iter()
                    .filter_map(|t| {
                        Some(McpTool {
                            server: cfg.name.clone(),
                            name: t["name"].as_str()?.to_string(),
                            description: t["description"].as_str().unwrap_or("").to_string(),
                            input_schema: t
                                .get("inputSchema")
                                .cloned()
                                .unwrap_or_else(|| json!({"type": "object", "properties": {}})),
                        })
                    })
                    .collect::<Vec<_>>()
            })
            .unwrap_or_default();

        Ok((conn, tools))
    }

    /// Call an MCP tool by its LLM-facing name (`mcp__server__tool`).
    pub async fn call(&mut self, llm_name: &str, arguments: Value) -> String {
        let Some(rest) = llm_name.strip_prefix("mcp__") else {
            return format!("Error: '{llm_name}' is not an MCP tool name.");
        };
        let Some((server, tool)) = rest.split_once("__") else {
            return format!("Error: malformed MCP tool name '{llm_name}'.");
        };
        let Some(conn) = self.connections.get_mut(server) else {
            return format!("Error: MCP server '{server}' is not connected.");
        };
        match conn
            .request("tools/call", json!({ "name": tool, "arguments": arguments }))
            .await
        {
            Ok(result) => {
                let mut out = String::new();
                if let Some(content) = result["content"].as_array() {
                    for item in content {
                        match item["type"].as_str() {
                            Some("text") => {
                                out.push_str(item["text"].as_str().unwrap_or(""));
                                out.push('\n');
                            }
                            Some(other) => {
                                out.push_str(&format!("[{other} content omitted]\n"));
                            }
                            None => {}
                        }
                    }
                }
                if result["isError"].as_bool().unwrap_or(false) {
                    out = format!("Tool reported an error:\n{out}");
                }
                if out.trim().is_empty() {
                    out = result.to_string();
                }
                crate::tools::shell::truncate_output(&out)
            }
            Err(e) => format!("MCP call failed: {e:#}"),
        }
    }

    /// OpenAI-style tool definitions for all connected MCP tools.
    pub fn tool_definitions(&self) -> Vec<Value> {
        self.tools
            .iter()
            .map(|t| {
                json!({
                    "type": "function",
                    "function": {
                        "name": t.llm_name(),
                        "description": format!("[MCP:{}] {}", t.server, t.description),
                        "parameters": t.input_schema,
                    }
                })
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn llm_name_roundtrip() {
        let t = McpTool {
            server: "github".into(),
            name: "search_issues".into(),
            description: String::new(),
            input_schema: json!({}),
        };
        assert_eq!(t.llm_name(), "mcp__github__search_issues");
    }

    #[tokio::test]
    async fn connects_to_a_scripted_server() {
        // A fake MCP server in bash: replies to initialize and tools/list.
        let script = r#"
while IFS= read -r line; do
  id=$(echo "$line" | sed -n 's/.*"id":\([0-9]*\).*/\1/p')
  case "$line" in
    *initialize*) echo "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{\"protocolVersion\":\"2024-11-05\",\"capabilities\":{},\"serverInfo\":{\"name\":\"fake\",\"version\":\"0\"}}}" ;;
    *tools/list*) echo "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{\"tools\":[{\"name\":\"echo\",\"description\":\"echoes\",\"inputSchema\":{\"type\":\"object\"}}]}}" ;;
    *tools/call*) echo "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"pong\"}]}}" ;;
  esac
done
"#;
        let cfg = McpServerConfig {
            name: "fake".into(),
            command: "bash".into(),
            args: vec!["-c".into(), script.into()],
            ..Default::default()
        };
        let mut mgr = McpManager::start(&[cfg]).await;
        assert!(mgr.errors.is_empty(), "{:?}", mgr.errors);
        assert_eq!(mgr.tools.len(), 1);
        assert_eq!(mgr.tools[0].llm_name(), "mcp__fake__echo");
        let out = mgr.call("mcp__fake__echo", json!({"msg": "hi"})).await;
        assert!(out.contains("pong"), "{out}");
    }
}
