//! MCP client shaped after Grok Bot: GetMcpTools then CallMcpTool.
//!
//! Config: `~/.grokboy/mcp.json` (override with `GROKBOY_MCP_CONFIG`).
//! Workspace `.grokboy/mcp.json` is merged on top.

use anyhow::{anyhow, Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::Mutex;

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct McpFile {
    #[serde(default, rename = "mcpServers")]
    pub mcp_servers: HashMap<String, McpServerConfig>,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct McpServerConfig {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub command: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub args: Vec<String>,
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub env: HashMap<String, String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub url: Option<String>,
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub headers: HashMap<String, String>,
}

impl McpServerConfig {
    fn transport(&self) -> &'static str {
        if self.url.as_ref().is_some_and(|u| !u.is_empty()) {
            "http"
        } else {
            "stdio"
        }
    }
}

pub fn mcp_config_path() -> Result<PathBuf> {
    if let Some(path) = std::env::var_os("GROKBOY_MCP_CONFIG").filter(|s| !s.is_empty()) {
        return Ok(PathBuf::from(path));
    }
    let home = std::env::var_os("HOME")
        .or_else(|| std::env::var_os("USERPROFILE"))
        .map(PathBuf::from)
        .ok_or_else(|| anyhow!("cannot resolve home directory"))?;
    Ok(home.join(".grokboy").join("mcp.json"))
}

pub fn load_mcp_file(path: &Path) -> Result<McpFile> {
    if !path.exists() {
        return Ok(McpFile::default());
    }
    let text = std::fs::read_to_string(path)
        .with_context(|| format!("read {}", path.display()))?;
    serde_json::from_str(&text).with_context(|| format!("parse {}", path.display()))
}

pub fn save_mcp_file(path: &Path, file: &McpFile) -> Result<()> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    std::fs::write(path, serde_json::to_vec_pretty(file)?)
        .with_context(|| format!("write {}", path.display()))
}

fn merge_workspace(cwd: &Path, mut file: McpFile) -> McpFile {
    let extra = cwd.join(".grokboy").join("mcp.json");
    if let Ok(local) = load_mcp_file(&extra) {
        file.mcp_servers.extend(local.mcp_servers);
    }
    file
}

struct StdioSession {
    #[allow(dead_code)]
    child: Child,
    stdin: ChildStdin,
    stdout: BufReader<ChildStdout>,
    next_id: u64,
}

enum Live {
    Stdio(StdioSession),
    Http {
        url: String,
        headers: HashMap<String, String>,
        next_id: u64,
    },
}

pub struct McpHub {
    path: PathBuf,
    file: Mutex<McpFile>,
    live: Mutex<HashMap<String, Live>>,
}

impl McpHub {
    pub fn empty() -> Arc<Self> {
        Arc::new(Self {
            path: std::env::temp_dir().join("grokboy-mcp-empty.json"),
            file: Mutex::new(McpFile::default()),
            live: Mutex::new(HashMap::new()),
        })
    }

    pub fn load(cwd: &Path) -> Result<Arc<Self>> {
        let path = mcp_config_path()?;
        let file = merge_workspace(cwd, load_mcp_file(&path)?);
        Ok(Arc::new(Self {
            path,
            file: Mutex::new(file),
            live: Mutex::new(HashMap::new()),
        }))
    }

    pub async fn status(&self, server: Option<&str>) -> Result<Value> {
        let file = self.file.lock().await;
        let mut rows = Vec::new();
        for (name, cfg) in file.mcp_servers.iter() {
            if server.is_some_and(|s| s != name) {
                continue;
            }
            rows.push(json!({
                "server": name,
                "status": "configured",
                "transport": cfg.transport(),
                "command": cfg.command,
                "url": cfg.url,
            }));
        }
        rows.sort_by(|a, b| a["server"].as_str().cmp(&b["server"].as_str()));
        Ok(json!({"servers": rows, "config": self.path}))
    }

    pub async fn add(
        &self,
        name: &str,
        url: Option<&str>,
        headers: Option<HashMap<String, String>>,
        command: Option<&str>,
        args: Option<Vec<String>>,
        env: Option<HashMap<String, String>>,
    ) -> Result<Value> {
        if name.trim().is_empty() {
            return Err(anyhow!("name is required"));
        }
        let mut cfg = McpServerConfig::default();
        if let Some(url) = url.filter(|s| !s.is_empty()) {
            if let Some(err) = validate_remote_url(url) {
                return Err(anyhow!(err));
            }
            cfg.url = Some(url.to_string());
            cfg.headers = headers.unwrap_or_default();
        } else if let Some(command) = command.filter(|s| !s.is_empty()) {
            cfg.command = Some(command.to_string());
            cfg.args = args.unwrap_or_default();
            cfg.env = env.unwrap_or_default();
        } else {
            return Err(anyhow!(
                "provide url (https MCP endpoint) or command (stdio server)"
            ));
        }
        let mut file = self.file.lock().await;
        file.mcp_servers.insert(name.to_string(), cfg);
        save_mcp_file(&self.path, &file)?;
        drop(file);
        self.live.lock().await.remove(name);
        self.status(Some(name)).await
    }

    pub async fn remove(&self, name: &str) -> Result<Value> {
        let mut file = self.file.lock().await;
        let removed = file.mcp_servers.remove(name).is_some();
        save_mcp_file(&self.path, &file)?;
        drop(file);
        self.live.lock().await.remove(name);
        Ok(json!({"removed": removed, "server": name}))
    }

    pub async fn get_tools(
        &self,
        server: Option<&str>,
        tool_name: Option<&str>,
        pattern: Option<&str>,
    ) -> Result<Value> {
        let names: Vec<String> = {
            let file = self.file.lock().await;
            let mut names: Vec<String> = file.mcp_servers.keys().cloned().collect();
            names.sort();
            if let Some(server) = server {
                if !file.mcp_servers.contains_key(server) {
                    return Err(anyhow!(
                        "MCP server `{server}` is not installed. GetMcpServerStatus lists identifiers; AddMcpServer installs one."
                    ));
                }
                names.retain(|n| n == server);
            }
            names
        };
        if names.is_empty() {
            return Ok(json!({
                "mode": "catalog",
                "servers": [],
                "instruction": "No MCP servers are installed. Use AddMcpServer with a stdio command or an https URL."
            }));
        }
        let re = pattern
            .filter(|s| !s.is_empty())
            .map(|p| regex::Regex::new(&format!("(?i){p}")))
            .transpose()
            .map_err(|e| anyhow!("invalid pattern: {e}"))?;
        let mut servers = Vec::new();
        for name in names {
            match self.list_server_tools(&name).await {
                Ok(mut tools) => {
                    if let Some(want) = tool_name.filter(|s| !s.is_empty()) {
                        tools.retain(|t| t["name"].as_str() == Some(want));
                    }
                    if let Some(re) = &re {
                        tools.retain(|t| {
                            let blob = format!(
                                "{} {}",
                                t["name"].as_str().unwrap_or(""),
                                t["description"].as_str().unwrap_or("")
                            );
                            re.is_match(&blob)
                        });
                    }
                    servers.push(json!({
                        "server": name,
                        "serverStatus": "connected",
                        "tools": tools,
                    }));
                }
                Err(err) => servers.push(json!({
                    "server": name,
                    "serverStatus": "error",
                    "serverError": err.to_string(),
                    "tools": [],
                })),
            }
        }
        Ok(json!({
            "mode": if tool_name.is_some() { "single_tool" } else { "catalog" },
            "servers": servers,
            "instruction": "Call a listed tool with CallMcpTool using the server identifier and tool name. Refetch this descriptor if a later call fails."
        }))
    }

    pub async fn call_tool(&self, server: &str, tool: &str, arguments: Value) -> Result<Value> {
        if server.trim().is_empty() || tool.trim().is_empty() {
            return Err(anyhow!("CallMcpTool requires server and tool_name"));
        }
        let exists = self.file.lock().await.mcp_servers.contains_key(server);
        if !exists {
            return Err(anyhow!(
                "MCP server `{server}` does not exist. GetMcpTools / GetMcpServerStatus list identifiers."
            ));
        }
        let raw = self.rpc(server, "tools/call", json!({"name": tool, "arguments": arguments})).await?;
        Ok(json!({
            "server": server,
            "tool": tool,
            "result": raw,
        }))
    }

    async fn list_server_tools(&self, server: &str) -> Result<Vec<Value>> {
        let raw = self.rpc(server, "tools/list", json!({})).await?;
        let tools = raw["tools"].as_array().cloned().unwrap_or_default();
        Ok(tools
            .into_iter()
            .map(|t| {
                json!({
                    "name": t["name"],
                    "description": t["description"].as_str().unwrap_or("").chars().take(200).collect::<String>(),
                    "inputSchema": t.get("inputSchema").cloned().unwrap_or(json!({"type":"object"})),
                })
            })
            .collect())
    }

    async fn rpc(&self, server: &str, method: &str, params: Value) -> Result<Value> {
        self.ensure(server).await?;
        let mut live = self.live.lock().await;
        let session = live
            .get_mut(server)
            .ok_or_else(|| anyhow!("MCP server `{server}` failed to start"))?;
        match session {
            Live::Stdio(stdio) => stdio.request(method, params).await,
            Live::Http {
                url,
                headers,
                next_id,
            } => http_request(url, headers, next_id, method, params).await,
        }
    }

    async fn ensure(&self, server: &str) -> Result<()> {
        if self.live.lock().await.contains_key(server) {
            return Ok(());
        }
        let cfg = self
            .file
            .lock()
            .await
            .mcp_servers
            .get(server)
            .cloned()
            .ok_or_else(|| anyhow!("unknown MCP server `{server}`"))?;
        let live = connect(&cfg).await?;
        self.live.lock().await.insert(server.to_string(), live);
        Ok(())
    }
}

fn validate_remote_url(raw: &str) -> Option<String> {
    let parsed = match reqwest::Url::parse(raw) {
        Ok(url) => url,
        Err(_) => return Some(format!("\"{raw}\" is not a valid URL.")),
    };
    if parsed.scheme() != "http" && parsed.scheme() != "https" {
        return Some(format!(
            "The server URL must be http(s); \"{}\" is not supported.",
            parsed.scheme()
        ));
    }
    if !parsed.username().is_empty() || parsed.password().is_some() {
        return Some(
            "Don't put credentials in the server URL — pass them as headers instead.".into(),
        );
    }
    None
}

async fn connect(cfg: &McpServerConfig) -> Result<Live> {
    if let Some(url) = cfg.url.as_ref().filter(|s| !s.is_empty()) {
        let mut live = Live::Http {
            url: url.clone(),
            headers: cfg.headers.clone(),
            next_id: 1,
        };
        initialize(&mut live).await?;
        return Ok(live);
    }
    let command = cfg
        .command
        .as_deref()
        .ok_or_else(|| anyhow!("stdio MCP server needs command"))?;
    let mut cmd = Command::new(command);
    cmd.args(&cfg.args)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .kill_on_drop(true);
    for (k, v) in &cfg.env {
        cmd.env(k, v);
    }
    let mut child = cmd.spawn().with_context(|| format!("spawn MCP `{command}`"))?;
    let stdin = child.stdin.take().ok_or_else(|| anyhow!("mcp stdin"))?;
    let stdout = BufReader::new(child.stdout.take().ok_or_else(|| anyhow!("mcp stdout"))?);
    let mut live = Live::Stdio(StdioSession {
        child,
        stdin,
        stdout,
        next_id: 1,
    });
    initialize(&mut live).await?;
    Ok(live)
}

async fn initialize(live: &mut Live) -> Result<()> {
    let result = match live {
        Live::Stdio(stdio) => {
            stdio
                .request(
                    "initialize",
                    json!({
                        "protocolVersion": "2024-11-05",
                        "capabilities": {},
                        "clientInfo": {"name": "grokboy", "version": "0.1.0"}
                    }),
                )
                .await?
        }
        Live::Http {
            url,
            headers,
            next_id,
        } => {
            http_request(
                url,
                headers,
                next_id,
                "initialize",
                json!({
                    "protocolVersion": "2024-11-05",
                    "capabilities": {},
                    "clientInfo": {"name": "grokboy", "version": "0.1.0"}
                }),
            )
            .await?
        }
    };
    let _ = result;
    match live {
        Live::Stdio(stdio) => {
            write_frame(
                &mut stdio.stdin,
                &json!({"jsonrpc":"2.0","method":"notifications/initialized"}),
            )
            .await
        }
        Live::Http { .. } => Ok(()),
    }
}

impl StdioSession {
    async fn request(&mut self, method: &str, params: Value) -> Result<Value> {
        let id = self.next_id;
        self.next_id += 1;
        write_frame(
            &mut self.stdin,
            &json!({"jsonrpc":"2.0","id":id,"method":method,"params":params}),
        )
        .await?;
        loop {
            let msg = read_frame(&mut self.stdout).await?;
            if msg.get("id") == Some(&json!(id)) {
                if let Some(err) = msg.get("error") {
                    return Err(anyhow!("MCP {method} error: {err}"));
                }
                return Ok(msg.get("result").cloned().unwrap_or(json!({})));
            }
        }
    }
}

async fn write_frame(stdin: &mut ChildStdin, value: &Value) -> Result<()> {
    let body = serde_json::to_vec(value)?;
    let header = format!("Content-Length: {}\r\n\r\n", body.len());
    stdin.write_all(header.as_bytes()).await?;
    stdin.write_all(&body).await?;
    stdin.flush().await?;
    Ok(())
}

async fn read_frame(stdout: &mut BufReader<ChildStdout>) -> Result<Value> {
    let mut content_length = None;
    loop {
        let mut line = String::new();
        let n = stdout.read_line(&mut line).await?;
        if n == 0 {
            return Err(anyhow!("MCP server closed stdout"));
        }
        if line == "\r\n" || line == "\n" {
            break;
        }
        if let Some(rest) = line
            .to_ascii_lowercase()
            .strip_prefix("content-length:")
            .or_else(|| line.strip_prefix("Content-Length:"))
        {
            content_length = rest.trim().parse().ok();
        }
    }
    let len = content_length.ok_or_else(|| anyhow!("MCP response missing Content-Length"))?;
    let mut buf = vec![0u8; len];
    stdout.read_exact(&mut buf).await?;
    serde_json::from_slice(&buf).context("parse MCP JSON-RPC")
}

async fn http_request(
    url: &str,
    headers: &HashMap<String, String>,
    next_id: &mut u64,
    method: &str,
    params: Value,
) -> Result<Value> {
    let id = *next_id;
    *next_id += 1;
    let mut req = reqwest::Client::new()
        .post(url)
        .header("content-type", "application/json")
        .json(&json!({"jsonrpc":"2.0","id":id,"method":method,"params":params}));
    for (k, v) in headers {
        req = req.header(k, v);
    }
    let parsed: Value = req.send().await?.json().await?;
    if let Some(err) = parsed.get("error") {
        return Err(anyhow!("MCP {method} error: {err}"));
    }
    Ok(parsed.get("result").cloned().unwrap_or(json!({})))
}

pub fn parse_headers(value: &Value) -> Option<HashMap<String, String>> {
    let obj = value.as_object()?;
    let mut out = HashMap::new();
    for (k, v) in obj {
        if let Some(s) = v.as_str() {
            out.insert(k.clone(), s.to_string());
        }
    }
    Some(out)
}

pub fn parse_env(value: &Value) -> Option<HashMap<String, String>> {
    parse_headers(value)
}

pub fn parse_args_list(value: &Value) -> Option<Vec<String>> {
    value
        .as_array()
        .map(|a| {
            a.iter()
                .filter_map(|v| v.as_str().map(str::to_string))
                .collect()
        })
}

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

    #[test]
    fn rejects_credentials_in_url() {
        assert!(validate_remote_url("https://user:pw@example.com/mcp").is_some());
        assert!(validate_remote_url("https://example.com/mcp").is_none());
    }

    #[test]
    fn round_trips_config() {
        let dir = std::env::temp_dir().join(format!("gb-mcp-{}", uuid::Uuid::new_v4()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("mcp.json");
        let mut file = McpFile::default();
        file.mcp_servers.insert(
            "echo".into(),
            McpServerConfig {
                command: Some("python3".into()),
                args: vec!["-c".into(), "pass".into()],
                ..Default::default()
            },
        );
        save_mcp_file(&path, &file).unwrap();
        let loaded = load_mcp_file(&path).unwrap();
        assert_eq!(loaded.mcp_servers["echo"].command.as_deref(), Some("python3"));
        let _ = std::fs::remove_dir_all(dir);
    }

    const ECHO_PY: &str = r#"
import json, sys
def read_msg():
    headers = {}
    while True:
        line = sys.stdin.buffer.readline()
        if not line or line in (b"\r\n", b"\n"):
            break
        k, _, v = line.decode().partition(":")
        headers[k.strip().lower()] = v.strip()
    n = int(headers.get("content-length", "0"))
    return json.loads(sys.stdin.buffer.read(n))
def write_msg(obj):
    data = json.dumps(obj).encode()
    sys.stdout.buffer.write(f"Content-Length: {len(data)}\r\n\r\n".encode() + data)
    sys.stdout.buffer.flush()
while True:
    msg = read_msg()
    method = msg.get("method")
    mid = msg.get("id")
    if method == "initialize":
        write_msg({"jsonrpc":"2.0","id":mid,"result":{"protocolVersion":"2024-11-05","capabilities":{"tools":{}},"serverInfo":{"name":"echo","version":"0"}}})
    elif method == "notifications/initialized":
        continue
    elif method == "tools/list":
        write_msg({"jsonrpc":"2.0","id":mid,"result":{"tools":[{"name":"echo","description":"echo text","inputSchema":{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}}]}})
    elif method == "tools/call":
        text = ((msg.get("params") or {}).get("arguments") or {}).get("text", "")
        write_msg({"jsonrpc":"2.0","id":mid,"result":{"content":[{"type":"text","text": text}]}})
"#;

    #[tokio::test]
    async fn stdio_get_mcp_tools_and_call() {
        let _lock = crate::test_env::lock_async().await;
        if std::process::Command::new("python3")
            .arg("-c")
            .arg("pass")
            .status()
            .ok()
            .is_none_or(|s| !s.success())
        {
            return;
        }
        let dir = std::env::temp_dir().join(format!("gb-mcp-live-{}", uuid::Uuid::new_v4()));
        std::fs::create_dir_all(&dir).unwrap();
        let script = dir.join("echo.py");
        std::fs::write(&script, ECHO_PY).unwrap();
        let cfg = dir.join("mcp.json");
        unsafe { std::env::set_var("GROKBOY_MCP_CONFIG", &cfg) };
        let mut file = McpFile::default();
        file.mcp_servers.insert(
            "echo".into(),
            McpServerConfig {
                command: Some("python3".into()),
                args: vec![script.to_string_lossy().into_owned()],
                ..Default::default()
            },
        );
        save_mcp_file(&cfg, &file).unwrap();
        let hub = McpHub::load(&dir).unwrap();
        let listed = hub.get_tools(Some("echo"), None, None).await.unwrap();
        assert_eq!(listed["servers"][0]["tools"][0]["name"], "echo");
        let called = hub
            .call_tool("echo", "echo", json!({"text": "hello-mcp"}))
            .await
            .unwrap();
        let blob = called.to_string();
        assert!(blob.contains("hello-mcp"), "{blob}");
        unsafe { std::env::remove_var("GROKBOY_MCP_CONFIG") };
        let _ = std::fs::remove_dir_all(dir);
    }
}