// MCP connectors: the bot can use the tools of Model Context Protocol servers (the original's
// local-exec-daemon has an MCP client with its own permission layer: shouldBlockMcp /
// isMcpFullyAllowlisted, and per-server custom instructions). Here: stdio servers (a command
// AnyBot starts) and HTTP servers (Streamable HTTP, JSON or SSE answers).

use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::{BTreeMap, HashMap};
use std::io::{BufRead, BufReader, Write};
use std::process::{Child, ChildStdin, Command, Stdio};
use std::sync::mpsc::{channel, Receiver};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;

pub const PROTOCOL: &str = "2025-06-18";
/// Longest wait for one answer (estimated).
pub const CALL_TIMEOUT_SECS: u64 = 120;
/// Tool names the model sees: mcp__<server>__<tool>, at most 64 characters.
pub const PREFIX: &str = "mcp__";

#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
pub struct Server {
    pub name: String,
    /// stdio: the command; empty when `url` is used.
    #[serde(default)]
    pub command: String,
    #[serde(default)]
    pub args: Vec<String>,
    #[serde(default)]
    pub env: BTreeMap<String, String>,
    /// HTTP (Streamable HTTP) address.
    #[serde(default)]
    pub url: String,
    /// e.g. "Authorization: Bearer …" for HTTP servers.
    #[serde(default)]
    pub headers: BTreeMap<String, String>,
    #[serde(default = "yes")]
    pub enabled: bool,
    /// Notes for the bot on how to use this server (the original's MCP custom instructions).
    #[serde(default)]
    pub instructions: String,
}

fn yes() -> bool {
    true
}

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

enum Link {
    Stdio { child: Child, stdin: ChildStdin, rx: Receiver<Value> },
    Http { url: String, headers: BTreeMap<String, String>, session: Option<String> },
}

/// The part a tool call mutates; locked per server, so one slow call (up to two minutes)
/// never blocks other servers, the UI's status line, or the app's own sync.
struct Live {
    link: Link,
    next_id: u64,
}

struct Conn {
    live: Mutex<Live>,
    tools: Vec<ToolInfo>,
    config: Server,
}

/// One lock for the table, one lock per server's link.
fn conns() -> &'static Mutex<HashMap<String, Arc<Conn>>> {
    static C: OnceLock<Mutex<HashMap<String, Arc<Conn>>>> = OnceLock::new();
    C.get_or_init(|| Mutex::new(HashMap::new()))
}

pub fn safe(s: &str) -> String {
    s.chars().map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' }).collect()
}

/// A server's part of tool names: its name when it's plain ASCII, else a short stable code.
pub fn server_key(server: &str) -> String {
    if !server.is_empty() && server.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == ' ') {
        return safe(server);
    }
    let h = server.bytes().fold(0xcbf2_9ce4_8422_2325u64, |h, b| (h ^ b as u64).wrapping_mul(0x100_0000_01b3));
    format!("s{:06x}", h & 0xff_ffff)
}

pub fn tool_name(server: &str, tool: &str) -> String {
    let n = format!("{PREFIX}{}__{}", server_key(server), safe(tool));
    n.chars().take(64).collect()
}

pub fn is_mcp_tool(name: &str) -> bool {
    name.starts_with(PREFIX)
}

fn send_stdio(stdin: &mut ChildStdin, msg: &Value) -> Result<(), String> {
    let mut line = msg.to_string();
    line.push('\n');
    stdin.write_all(line.as_bytes()).and_then(|_| stdin.flush()).map_err(|e| format!("送不出去：{e}"))
}

/// Read the answer with this id (skipping notifications and requests from the server).
fn wait_id(rx: &Receiver<Value>, id: u64, secs: u64) -> Result<Value, String> {
    let deadline = std::time::Instant::now() + Duration::from_secs(secs);
    loop {
        let left = deadline.saturating_duration_since(std::time::Instant::now());
        if left.is_zero() {
            return Err("MCP 伺服器沒有回應".into());
        }
        match rx.recv_timeout(left) {
            Ok(v) if v["id"].as_u64() == Some(id) && v.get("method").is_none() => {
                if let Some(e) = v.get("error") {
                    return Err(format!("MCP 錯誤：{}", e["message"].as_str().unwrap_or(&e.to_string())));
                }
                return Ok(v["result"].clone());
            }
            Ok(_) => continue,
            Err(_) => return Err("MCP 伺服器沒有回應或已經結束".into()),
        }
    }
}

/// Answers from Streamable HTTP: plain JSON, or SSE with `data:` lines.
pub fn parse_http_answer(body: &str, id: u64) -> Result<Value, String> {
    let candidates: Vec<Value> = if body.trim_start().starts_with('{') || body.trim_start().starts_with('[') {
        match serde_json::from_str::<Value>(body) {
            Ok(Value::Array(a)) => a,
            Ok(v) => vec![v],
            Err(e) => return Err(e.to_string()),
        }
    } else {
        body.lines().filter_map(|l| l.strip_prefix("data:")).filter_map(|d| serde_json::from_str(d.trim()).ok()).collect()
    };
    for v in candidates {
        if v["id"].as_u64() == Some(id) {
            if let Some(e) = v.get("error") {
                return Err(format!("MCP 錯誤：{}", e["message"].as_str().unwrap_or(&e.to_string())));
            }
            return Ok(v["result"].clone());
        }
    }
    Err("MCP 伺服器沒有給答案".into())
}

impl Live {
    fn request(&mut self, method: &str, params: Value) -> Result<Value, String> {
        let id = self.next_id;
        self.next_id += 1;
        let msg = json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params});
        match &mut self.link {
            Link::Stdio { stdin, rx, .. } => {
                send_stdio(stdin, &msg)?;
                wait_id(rx, id, CALL_TIMEOUT_SECS)
            }
            Link::Http { url, headers, session } => {
                let agent: ureq::Agent = ureq::Agent::config_builder().timeout_global(Some(Duration::from_secs(CALL_TIMEOUT_SECS))).http_status_as_error(false).build().into();
                let mut r = agent.post(url.as_str()).header("Content-Type", "application/json").header("Accept", "application/json, text/event-stream").header("MCP-Protocol-Version", PROTOCOL);
                for (k, v) in headers.iter() {
                    r = r.header(k.as_str(), v.as_str());
                }
                if let Some(s) = session.as_ref() {
                    r = r.header("Mcp-Session-Id", s.as_str());
                }
                let mut resp = r.send(msg.to_string().as_bytes()).map_err(|e| format!("連不上 MCP 伺服器：{e}"))?;
                if let Some(s) = resp.headers().get("mcp-session-id").and_then(|v| v.to_str().ok()) {
                    *session = Some(s.to_string());
                }
                let st = resp.status().as_u16();
                let body = resp.body_mut().read_to_string().unwrap_or_default();
                if !(200..300).contains(&st) {
                    return Err(format!("MCP 伺服器回 HTTP {st}：{}", body.chars().take(200).collect::<String>()));
                }
                parse_http_answer(&body, id)
            }
        }
    }

    fn notify(&mut self, method: &str) {
        let msg = json!({"jsonrpc": "2.0", "method": method});
        match &mut self.link {
            Link::Stdio { stdin, .. } => {
                let _ = send_stdio(stdin, &msg);
            }
            Link::Http { url, headers, session } => {
                let mut r = ureq::post(url.as_str()).header("Content-Type", "application/json").header("Accept", "application/json, text/event-stream");
                for (k, v) in headers.iter() {
                    r = r.header(k.as_str(), v.as_str());
                }
                if let Some(s) = session.as_ref() {
                    r = r.header("Mcp-Session-Id", s.as_str());
                }
                let _ = r.send(msg.to_string().as_bytes());
            }
        }
    }
}

fn connect(s: &Server) -> Result<Conn, String> {
    let link = if !s.url.trim().is_empty() {
        Link::Http { url: s.url.trim().to_string(), headers: s.headers.clone(), session: None }
    } else {
        if s.command.trim().is_empty() {
            return Err("請給指令或網址".into());
        }
        let mut child = Command::new(s.command.trim())
            .args(&s.args)
            .envs(&s.env)
            .env("PATH", crate::cloudbox::tool_path())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .spawn()
            .map_err(|e| format!("無法啟動「{}」：{e}", s.command))?;
        let stdin = child.stdin.take().ok_or("沒有標準輸入")?;
        let stdout = child.stdout.take().ok_or("沒有標準輸出")?;
        let (tx, rx) = channel();
        std::thread::spawn(move || {
            for line in BufReader::new(stdout).lines().map_while(Result::ok) {
                if let Ok(v) = serde_json::from_str::<Value>(&line) {
                    if tx.send(v).is_err() {
                        break;
                    }
                }
            }
        });
        Link::Stdio { child, stdin, rx }
    };
    let mut live = Live { link, next_id: 1 };
    live.request("initialize", json!({"protocolVersion": PROTOCOL, "capabilities": {}, "clientInfo": {"name": "AnyBot", "version": "1"}}))?;
    live.notify("notifications/initialized");
    let mut tools = vec![];
    let mut cursor: Option<String> = None;
    for _ in 0..20 {
        let r = live.request("tools/list", match &cursor { Some(x) => json!({"cursor": x}), None => json!({}) })?;
        for t in r["tools"].as_array().into_iter().flatten() {
            tools.push(ToolInfo { server: s.name.clone(), name: t["name"].as_str().unwrap_or("").into(), description: t["description"].as_str().unwrap_or("").into(), schema: t["inputSchema"].clone() });
        }
        cursor = r["nextCursor"].as_str().map(String::from);
        if cursor.is_none() {
            break;
        }
    }
    Ok(Conn { live: Mutex::new(live), tools, config: s.clone() })
}

/// Start (or keep) the enabled servers; stop the rest. Returns errors per server.
/// Connecting happens outside the table's lock (a slow server must not freeze the UI's
/// status line, which reads that lock every second).
pub fn sync(servers: &[Server]) -> Vec<(String, String)> {
    let mut errs = vec![];
    let mut to_connect: Vec<Server> = vec![];
    {
        let mut map = conns().lock().unwrap();
        let wanted: Vec<&Server> = servers.iter().filter(|s| s.enabled).collect();
        map.retain(|name, c| {
            let keep = wanted.iter().any(|s| &s.name == name && **s == c.config);
            if !keep {
                if let Link::Stdio { child, .. } = &mut c.live.lock().unwrap().link {
                    let _ = child.kill();
                }
            }
            keep
        });
        for s in wanted {
            if !map.contains_key(&s.name) {
                to_connect.push(s.clone());
            }
        }
    }
    for s in to_connect {
        match connect(&s) {
            Ok(c) => {
                conns().lock().unwrap().insert(s.name.clone(), Arc::new(c));
            }
            Err(e) => errs.push((s.name.clone(), e)),
        }
    }
    errs
}

/// The connected servers' tools.
pub fn tools() -> Vec<ToolInfo> {
    conns().lock().unwrap().values().flat_map(|c| c.tools.clone()).collect()
}

pub fn status() -> HashMap<String, usize> {
    conns().lock().unwrap().iter().map(|(k, c)| (k.clone(), c.tools.len())).collect()
}

/// Tool specs for the model.
pub fn specs() -> Vec<Value> {
    tools()
        .into_iter()
        .map(|t| {
            let schema = if t.schema.is_object() { t.schema.clone() } else { json!({"type": "object", "properties": {}}) };
            json!({"name": tool_name(&t.server, &t.name), "description": format!("[{}] {}", t.server, t.description.chars().take(900).collect::<String>()), "parameters": schema})
        })
        .collect()
}

/// Instructions the servers' owners wrote for the bot.
pub fn instructions(servers: &[Server]) -> String {
    servers.iter().filter(|s| s.enabled && !s.instructions.trim().is_empty()).map(|s| format!("- {}：{}", s.name, s.instructions.trim())).collect::<Vec<_>>().join("\n")
}

/// Which server and tool a model-side name means.
pub fn resolve(name: &str) -> Option<(String, String)> {
    tools().into_iter().find(|t| tool_name(&t.server, &t.name) == name).map(|t| (t.server, t.name))
}

/// Call a tool. Returns (text, ok, first picture as base64).
pub fn call(name: &str, args: &str) -> (String, bool, Option<String>) {
    let Some((server, tool)) = resolve(name) else { return (format!("沒有「{name}」這個 MCP 工具（伺服器可能已關閉）"), false, None) };
    let arguments: Value = serde_json::from_str(if args.trim().is_empty() { "{}" } else { args }).unwrap_or(json!({}));
    // Copy the server's own handle out of the table, then lock its link alone: a two-minute
    // call to one server used to block every other MCP tool and the UI's status line.
    let conn = conns().lock().unwrap().get(&server).cloned();
    let Some(conn) = conn else { return ("MCP 伺服器沒有連線".into(), false, None) };
    let mut live = conn.live.lock().unwrap();
    match live.request("tools/call", json!({"name": tool, "arguments": arguments})) {
        Ok(r) => {
            let mut text = String::new();
            let mut image = None;
            for part in r["content"].as_array().into_iter().flatten() {
                match part["type"].as_str() {
                    Some("text") => {
                        text.push_str(part["text"].as_str().unwrap_or(""));
                        text.push('\n');
                    }
                    Some("image") if image.is_none() => image = part["data"].as_str().map(String::from),
                    Some("resource") => {
                        text.push_str(part["resource"]["text"].as_str().unwrap_or(""));
                        text.push('\n');
                    }
                    _ => {}
                }
            }
            if text.trim().is_empty() && r.get("structuredContent").is_some() {
                text = r["structuredContent"].to_string();
            }
            let ok = r["isError"].as_bool() != Some(true);
            (crate::tools::head_tail(text.trim()), ok, image)
        }
        Err(e) => (e, false, None),
    }
}

pub fn shutdown() {
    for c in conns().lock().unwrap().values() {
        if let Link::Stdio { child, .. } = &mut c.live.lock().unwrap().link {
            let _ = child.kill();
        }
    }
}

/// 連接器市集: (name, what it does, command or address, key to fill in — empty when none).
pub const MARKET: [(&str, &str, &str, &str); 8] = [
    ("files", "讀寫你指定資料夾裡的檔案", "npx -y @modelcontextprotocol/server-filesystem ~/Documents", ""),
    ("memory", "長期記憶：把人、事、關係記成知識圖", "npx -y @modelcontextprotocol/server-memory", ""),
    ("thinking", "一步步拆解難題", "npx -y @modelcontextprotocol/server-sequential-thinking", ""),
    ("fetch", "抓網頁內容轉成文字", "uvx mcp-server-fetch", ""),
    ("git", "查看和操作 Git 版本庫", "uvx mcp-server-git", ""),
    ("time", "時間和時區換算", "uvx mcp-server-time", ""),
    ("playwright", "用瀏覽器自動化操作網頁", "npx -y @playwright/mcp@latest", ""),
    ("github", "GitHub 的 Issue、PR、程式碼（官方遠端伺服器）", "https://api.githubcopilot.com/mcp/", "Authorization: Bearer 你的 GitHub Token"),
];

/// A marketplace entry as a server (home folder spelled out).
pub fn market_server(name: &str, cmd: &str) -> Server {
    let parts: Vec<String> = cmd.split_whitespace().map(|w| if w == "~/Documents" { crate::tools::expand_home(w).display().to_string() } else { w.to_string() }).collect();
    Server { name: name.to_string(), command: parts[0].clone(), args: parts[1..].to_vec(), ..serde_json::from_str("{\"name\":\"x\"}").unwrap() }
}

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

    const FAKE_SERVER: &str = r#"
import sys, json
for line in sys.stdin:
    m = json.loads(line)
    if "id" not in m:
        continue
    if m["method"] == "initialize":
        r = {"protocolVersion": "2025-06-18", "capabilities": {"tools": {}}, "serverInfo": {"name": "fake"}}
    elif m["method"] == "tools/list":
        r = {"tools": [{"name": "add", "description": "Add two numbers", "inputSchema": {"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}}]}
    elif m["method"] == "tools/call":
        a = m["params"]["arguments"]
        r = {"content": [{"type": "text", "text": "sum=%d" % (a["a"] + a["b"])}]}
    else:
        r = {}
    sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": m["id"], "result": r}) + "\n")
    sys.stdout.flush()
"#;

    #[test]
    fn stdio_server_round_trip() {
        let dir = std::env::temp_dir().join(format!("anybot-mcp-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let script = dir.join("server.py");
        std::fs::write(&script, FAKE_SERVER).unwrap();
        let s = Server { name: "計算機".into(), command: "python3".into(), args: vec!["-I".into(), script.display().to_string()], env: BTreeMap::new(), url: String::new(), headers: BTreeMap::new(), enabled: true, instructions: "用來算數學".into() };
        let errs = sync(std::slice::from_ref(&s));
        assert!(errs.is_empty(), "{errs:?}");
        let name = tool_name("計算機", "add");
        assert!(name.starts_with("mcp__s") && name.ends_with("__add"), "{name}");
        assert_eq!(tool_name("github", "list issues"), "mcp__github__list_issues");
        assert!(specs().iter().any(|t| t["name"] == name.as_str()));
        let (out, ok, _) = call(&name, r#"{"a": 2, "b": 40}"#);
        assert!(ok && out.contains("sum=42"), "{out}");
        assert!(instructions(std::slice::from_ref(&s)).contains("用來算數學"));
        sync(&[]);
        assert!(tools().is_empty());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn http_answers_json_or_sse() {
        assert_eq!(parse_http_answer(r#"{"jsonrpc":"2.0","id":3,"result":{"ok":1}}"#, 3).unwrap()["ok"], 1);
        let sse = "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":4,\"result\":{\"tools\":[]}}\n\n";
        assert!(parse_http_answer(sse, 4).unwrap()["tools"].is_array());
        assert!(parse_http_answer(r#"{"jsonrpc":"2.0","id":5,"error":{"message":"nope"}}"#, 5).unwrap_err().contains("nope"));
    }

    /// Downloads and starts every marketplace connector that needs no key (needs the network,
    /// Node.js and uv): cargo test --release -- --ignored marketplace --nocapture
    #[test]
    #[ignore]
    fn marketplace_connectors_start() {
        let servers: Vec<Server> = MARKET.iter().filter(|m| m.3.is_empty()).map(|m| market_server(m.0, m.2)).collect();
        let errors = sync(&servers);
        let st = status();
        for s in &servers {
            println!("{:<12} {:?}", s.name, st.get(&s.name));
        }
        assert!(errors.is_empty(), "{errors:?}");
        for s in &servers {
            assert!(st.get(&s.name).copied().unwrap_or(0) > 0, "{} has no tools", s.name);
        }
        // And one real call through the bot's tool name.
        let (out, ok, _) = call(&tool_name("time", "get_current_time"), r#"{"timezone":"Asia/Taipei"}"#);
        println!("time: {out}");
        assert!(ok && out.contains("Asia/Taipei"), "{out}");
        shutdown();
    }
}