collet 0.1.0

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use serde::Deserialize;

// ---------------------------------------------------------------------------
// collet mcp.json schema
// ---------------------------------------------------------------------------
//
// Global format (~/.collet/mcp.json):
// {
//   "version": 1,
//   "mcpServers": {
//     "<name>": {
//       "source":      "npm:@pkg@ver | local:/path | http:https://... | github:owner/repo@tag",
//       "command":     "npx" | ["npx", "-y", "pkg"],   // stdio transport
//       "args":        ["-y", "@pkg"],
//       "env":         { "KEY": "${VAR}" },
//       "url":         "https://...",                   // HTTP transport
//       "headers":     { "Authorization": "Bearer ${TOKEN}" },
//       "enabled":     true,
//       "description": "human-readable description"
//     }
//   }
// }
//
// ---------------------------------------------------------------------------

/// Top-level shape of `mcp.json`.
///
/// Superseded by manual per-entry parsing in [`read_config_file_with_fixup`] which
/// handles command arrays and per-entry fixup. Kept for schema documentation and
/// validation in tests.
#[cfg(test)]
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct McpConfigFile {
    /// Schema version — reserved for future migrations.
    #[serde(default)]
    version: u32,
    #[serde(default)]
    mcp_servers: HashMap<String, McpServerEntry>,
    /// Legacy fallback key.
    #[serde(default)]
    servers: HashMap<String, McpServerEntry>,
}

/// A single MCP server entry.
///
/// Transport is inferred:
/// - `command` present → stdio
/// - `url` present     → HTTP
#[derive(Debug, Clone, Deserialize)]
pub struct McpServerEntry {
    // ── Origin ──────────────────────────────────────────────────────────────
    /// Install source — used by `collet mcp install` (future).
    /// Schemes: `npm:`, `local:`, `http:`, `github:`, `collet:`
    pub source: Option<String>,

    /// Human-readable description shown in `/mcp`.
    pub description: Option<String>,

    /// When false, server is skipped without removing the entry.
    #[serde(default = "default_true")]
    pub enabled: bool,

    // ── stdio transport ──────────────────────────────────────────────────────
    /// Executable to spawn. String or array (first element = binary).
    #[serde(default, deserialize_with = "deserialize_command")]
    pub command: Option<String>,
    /// Extra args extracted from a command array.
    #[serde(skip)]
    pub command_extra_args: Vec<String>,
    /// Arguments passed to the command.
    #[serde(default)]
    pub args: Vec<String>,
    /// Environment variables. Values support `${VAR}` expansion.
    #[serde(default)]
    pub env: HashMap<String, String>,
    /// Alias for `env`.
    #[serde(default)]
    pub environment: HashMap<String, String>,

    // ── HTTP transport ───────────────────────────────────────────────────────
    /// Base URL for HTTP MCP servers.
    pub url: Option<String>,
    /// HTTP headers (e.g. `Authorization`). Values support `${VAR}` expansion.
    #[serde(default)]
    pub headers: HashMap<String, String>,

    // ── ignored legacy fields ────────────────────────────────────────────────
    #[serde(rename = "type")]
    pub _type: Option<String>,
}

fn default_true() -> bool {
    true
}

/// Resolved server config ready for `McpClient` construction.
#[derive(Debug, Clone)]
pub struct ResolvedMcpServer {
    pub name: String,
    /// Install source hint (e.g. `npm:@playwright/mcp@latest`).
    pub source: Option<String>,
    /// Human-readable description logged during server connection.
    pub description: Option<String>,
    pub command: Option<String>,
    pub args: Vec<String>,
    /// Resolved environment variables (all `${VAR}` expanded).
    pub env: HashMap<String, String>,
    pub url: Option<String>,
    /// Resolved HTTP headers (all `${VAR}` expanded).
    pub headers: HashMap<String, String>,
}

/// MCP server availability status for sidebar display.
#[derive(Debug, Clone, PartialEq)]
pub enum McpStatus {
    /// Stdio server: binary found on $PATH.  HTTP server: URL configured.
    Available,
    /// Stdio server: binary NOT found on $PATH.
    Unavailable,
}

/// Entry for UiState.mcp_servers.
#[derive(Debug, Clone)]
pub struct McpServerUi {
    pub name: String,
    pub status: McpStatus,
}

// ---------------------------------------------------------------------------
// Loading
// ---------------------------------------------------------------------------

/// Search paths for `mcp.json` (project-local takes precedence over global).
///
/// - Project: `<cwd>/.collet/mcp.json`
/// - Global:  `~/.collet/mcp.json`
pub fn config_candidates(working_dir: &str) -> Vec<PathBuf> {
    vec![
        Path::new(working_dir).join(".collet").join("mcp.json"),
        crate::config::collet_home(None).join("mcp.json"),
    ]
}

/// Load and merge MCP server configs from all candidate paths.
/// Project-local entries take precedence over user-global ones.
pub fn load_mcp_servers(working_dir: &str) -> Result<Vec<ResolvedMcpServer>> {
    let mut merged: HashMap<String, McpServerEntry> = HashMap::new();

    // User-global first (lower priority) …
    for path in config_candidates(working_dir).into_iter().rev() {
        if let Some(entries) = read_config_file_with_fixup(&path)? {
            for (name, entry) in entries {
                merged.insert(name, entry);
            }
        }
    }

    let servers = merged
        .into_iter()
        .filter(|(_, entry)| entry.enabled)
        .map(|(name, entry)| resolve_entry(name, entry))
        .collect();
    Ok(servers)
}

/// Load MCP servers and check availability for sidebar display.
pub fn load_mcp_status(working_dir: &str) -> Vec<McpServerUi> {
    let servers = load_mcp_servers(working_dir).unwrap_or_default();
    let mut result: Vec<McpServerUi> = servers
        .into_iter()
        .map(|s| {
            let status = check_availability(&s);
            McpServerUi {
                name: s.name,
                status,
            }
        })
        .collect();
    result.sort_by(|a, b| a.name.cmp(&b.name));
    result
}

/// Check if a resolved MCP server is reachable.
fn check_availability(server: &ResolvedMcpServer) -> McpStatus {
    if let Some(ref cmd) = server.command {
        if is_binary_available(cmd) {
            McpStatus::Available
        } else {
            McpStatus::Unavailable
        }
    } else if server.url.is_some() {
        // HTTP servers: assume available if URL is configured.
        // Real connectivity check would require async; deferred to runtime.
        McpStatus::Available
    } else {
        McpStatus::Unavailable
    }
}

/// Check if a binary is available on $PATH.
fn is_binary_available(command: &str) -> bool {
    std::process::Command::new("which")
        .arg(command)
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status()
        .map(|s| s.success())
        .unwrap_or(false)
}

/// Expand `${VAR}` references in env values and merge `env` + `environment`.
fn resolve_entry(name: String, entry: McpServerEntry) -> ResolvedMcpServer {
    // Merge environment (alias) into env; env takes precedence.
    let mut env_map = entry.environment;
    env_map.extend(entry.env);
    let env = env_map
        .into_iter()
        .map(|(k, v)| (k, expand_env_vars(&v)))
        .collect();

    // Merge command_extra_args in front of args.
    let mut args = entry.command_extra_args;
    args.extend(entry.args);

    let headers = entry
        .headers
        .into_iter()
        .map(|(k, v)| (k, expand_env_vars(&v)))
        .collect();

    ResolvedMcpServer {
        name,
        source: entry.source,
        description: entry.description,
        command: entry.command,
        args,
        env,
        url: entry.url,
        headers,
    }
}

/// Replace `${VAR_NAME}` patterns with the corresponding environment variable.
/// If the variable is not set, the placeholder is replaced with an empty string.
fn expand_env_vars(input: &str) -> String {
    let mut result = String::with_capacity(input.len());
    let mut chars = input.chars().peekable();

    while let Some(ch) = chars.next() {
        if ch == '$' && chars.peek() == Some(&'{') {
            chars.next(); // consume '{'
            let var_name: String = chars.by_ref().take_while(|&c| c != '}').collect();
            if !var_name.is_empty() {
                match std::env::var(&var_name) {
                    Ok(val) => result.push_str(&val),
                    Err(_) => {
                        tracing::warn!(
                            var = %var_name,
                            "MCP config references env var ${{{var_name}}} but it is not set — \
                             the value will be empty, which may cause MCP authentication failures"
                        );
                    }
                }
            }
        } else {
            result.push(ch);
        }
    }

    result
}

// ---------------------------------------------------------------------------
// Custom deserializer: command as string OR array
// ---------------------------------------------------------------------------

/// Deserialize `command` that can be either:
/// - A string: `"npx"` → command = "npx", extra_args = []
/// - An array: `["npx", "-y", "pkg"]` → command = "npx", extra_args = ["-y", "pkg"]
///
/// The extra_args are stored in a thread-local and merged in during entry
/// construction since serde doesn't let us write to a sibling field.
fn deserialize_command<'de, D>(deserializer: D) -> std::result::Result<Option<String>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    use serde::de;

    struct CommandVisitor;

    impl<'de> de::Visitor<'de> for CommandVisitor {
        type Value = Option<String>;

        fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
            f.write_str("a string or array of strings")
        }

        fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
            Ok(None)
        }

        fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
            Ok(None)
        }

        fn visit_str<E: de::Error>(self, v: &str) -> std::result::Result<Self::Value, E> {
            Ok(Some(v.to_string()))
        }

        fn visit_string<E: de::Error>(self, v: String) -> std::result::Result<Self::Value, E> {
            Ok(Some(v))
        }

        fn visit_seq<A: de::SeqAccess<'de>>(
            self,
            mut seq: A,
        ) -> std::result::Result<Self::Value, A::Error> {
            let first: Option<String> = seq.next_element()?;
            let cmd = match first {
                Some(s) => s,
                None => return Ok(None),
            };
            // Collect remaining elements as extra args in thread-local storage
            let mut extras = Vec::new();
            while let Some(arg) = seq.next_element::<String>()? {
                extras.push(arg);
            }
            if !extras.is_empty() {
                COMMAND_EXTRA_ARGS.with(|cell| {
                    *cell.borrow_mut() = extras;
                });
            }
            Ok(Some(cmd))
        }
    }

    deserializer.deserialize_any(CommandVisitor)
}

thread_local! {
    static COMMAND_EXTRA_ARGS: std::cell::RefCell<Vec<String>> = const { std::cell::RefCell::new(Vec::new()) };
}

/// Post-process a deserialized `McpServerEntry` to move thread-local extra args.
fn take_command_extra_args() -> Vec<String> {
    COMMAND_EXTRA_ARGS.with(|cell| std::mem::take(&mut *cell.borrow_mut()))
}

/// Read config and properly handle command arrays.
fn read_config_file_with_fixup(path: &Path) -> Result<Option<HashMap<String, McpServerEntry>>> {
    if !path.exists() {
        return Ok(None);
    }
    let raw = std::fs::read_to_string(path)
        .with_context(|| format!("Failed to read MCP config: {}", path.display()))?;

    // Parse each entry individually to capture thread-local state per entry.
    let val: serde_json::Value = serde_json::from_str(&raw)
        .with_context(|| format!("Failed to parse MCP config: {}", path.display()))?;

    let mut entries = HashMap::new();

    // Try mcpServers first, then servers
    for key in &["mcpServers", "servers"] {
        if let Some(serde_json::Value::Object(map)) = val.get(key) {
            for (name, entry_val) in map {
                // Clear thread-local before each entry
                COMMAND_EXTRA_ARGS.with(|cell| cell.borrow_mut().clear());
                if let Ok(mut entry) = serde_json::from_value::<McpServerEntry>(entry_val.clone()) {
                    entry.command_extra_args = take_command_extra_args();
                    entries.insert(name.clone(), entry);
                }
            }
        }
    }

    if entries.is_empty() {
        Ok(None)
    } else {
        Ok(Some(entries))
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn test_mcp_config_file_version_field() {
        // McpConfigFile is superseded by manual per-entry parsing, but remains the
        // canonical schema doc. Verify it deserializes the version and key fields.
        let json = r#"{"version": 2, "mcpServers": {}, "servers": {}}"#;
        let parsed: McpConfigFile = serde_json::from_str(json).unwrap();
        assert_eq!(parsed.version, 2);
        assert!(parsed.mcp_servers.is_empty());
        assert!(parsed.servers.is_empty());
    }

    #[test]
    fn test_expand_env_vars_simple() {
        // SAFETY: test-only, single-threaded test runner
        unsafe { std::env::set_var("_COLLET_TEST_KEY", "secret123") };
        assert_eq!(expand_env_vars("${_COLLET_TEST_KEY}"), "secret123");
        unsafe { std::env::remove_var("_COLLET_TEST_KEY") };
    }

    #[test]
    fn test_expand_env_vars_multiple() {
        unsafe { std::env::set_var("_COLLET_A", "hello") };
        unsafe { std::env::set_var("_COLLET_B", "world") };
        assert_eq!(expand_env_vars("${_COLLET_A}-${_COLLET_B}"), "hello-world");
        unsafe { std::env::remove_var("_COLLET_A") };
        unsafe { std::env::remove_var("_COLLET_B") };
    }

    #[test]
    fn test_expand_env_vars_missing() {
        assert_eq!(expand_env_vars("${_COLLET_NONEXISTENT_VAR}"), "");
    }

    #[test]
    fn test_expand_env_vars_no_placeholder() {
        assert_eq!(expand_env_vars("plain-value"), "plain-value");
    }

    #[test]
    fn test_expand_env_vars_partial_dollar() {
        assert_eq!(expand_env_vars("cost $5"), "cost $5");
    }

    #[test]
    fn test_parse_command_string() {
        let json = r#"{ "mcpServers": { "s1": { "command": "npx", "args": ["-y", "pkg"] } } }"#;
        let entries = parse_test_config(json);
        let s = &entries["s1"];
        assert_eq!(s.command.as_deref(), Some("npx"));
        assert!(s.command_extra_args.is_empty());
        assert_eq!(s.args, vec!["-y", "pkg"]);
    }

    #[test]
    fn test_parse_command_array() {
        let json = r#"{ "mcpServers": { "s1": { "command": ["npx", "-y", "pkg"] } } }"#;
        let entries = parse_test_config(json);
        let s = &entries["s1"];
        assert_eq!(s.command.as_deref(), Some("npx"));
        assert_eq!(s.command_extra_args, vec!["-y", "pkg"]);
    }

    #[test]
    fn test_parse_environment_alias() {
        let json = r#"{ "mcpServers": { "s1": {
            "command": "bin",
            "environment": { "FOO": "bar" }
        } } }"#;
        let entries = parse_test_config(json);
        assert_eq!(entries["s1"].environment["FOO"], "bar");

        let resolved = resolve_entry("s1".into(), entries.into_values().next().unwrap());
        assert_eq!(resolved.env["FOO"], "bar");
    }

    #[test]
    fn test_parse_headers() {
        let json = r#"{ "mcpServers": { "s1": {
            "url": "https://api.example.com/mcp",
            "headers": { "Authorization": "Bearer tok" }
        } } }"#;
        let entries = parse_test_config(json);
        let resolved = resolve_entry("s1".into(), entries.into_values().next().unwrap());
        assert_eq!(resolved.headers["Authorization"], "Bearer tok");
    }

    #[test]
    fn test_parse_type_and_enabled_ignored() {
        let json = r#"{ "mcpServers": { "s1": {
            "command": "bin",
            "type": "local",
            "enabled": true
        } } }"#;
        let entries = parse_test_config(json);
        assert_eq!(entries["s1"].command.as_deref(), Some("bin"));
    }

    #[test]
    fn test_parse_real_world_config() {
        let json = r#"{
            "mcpServers": {
                "alcove": {
                    "command": ["/usr/bin/alcove"],
                    "environment": { "DOCS_ROOT": "/tmp/docs" },
                    "type": "local"
                },
                "web-reader": {
                    "headers": { "Authorization": "Bearer ${_COLLET_TEST_WR}" },
                    "type": "remote",
                    "url": "https://api.example.com/mcp"
                },
                "playwright": {
                    "command": "npx",
                    "args": ["-y", "@playwright/mcp@latest"]
                }
            }
        }"#;

        unsafe { std::env::set_var("_COLLET_TEST_WR", "secret") };
        let entries = parse_test_config(json);
        assert_eq!(entries.len(), 3);

        // alcove: command array
        let alcove = &entries["alcove"];
        assert_eq!(alcove.command.as_deref(), Some("/usr/bin/alcove"));
        assert_eq!(alcove.environment["DOCS_ROOT"], "/tmp/docs");

        // web-reader: HTTP with headers
        let wr = &entries["web-reader"];
        assert!(wr.command.is_none());
        assert_eq!(wr.url.as_deref(), Some("https://api.example.com/mcp"));
        let resolved = resolve_entry("web-reader".into(), wr.clone());
        assert_eq!(resolved.headers["Authorization"], "Bearer secret");

        // playwright: normal string command
        let pw = &entries["playwright"];
        assert_eq!(pw.command.as_deref(), Some("npx"));
        assert_eq!(pw.args, vec!["-y", "@playwright/mcp@latest"]);

        unsafe { std::env::remove_var("_COLLET_TEST_WR") };
    }

    #[test]
    fn test_check_availability_http() {
        let server = ResolvedMcpServer {
            name: "test".into(),
            source: None,
            description: None,
            command: None,
            args: vec![],
            env: HashMap::new(),
            url: Some("https://example.com".into()),
            headers: HashMap::new(),
        };
        assert_eq!(check_availability(&server), McpStatus::Available);
    }

    #[test]
    fn test_check_availability_missing_binary() {
        let server = ResolvedMcpServer {
            name: "test".into(),
            source: None,
            description: None,
            command: Some("__collet_nonexistent_binary__".into()),
            args: vec![],
            env: HashMap::new(),
            url: None,
            headers: HashMap::new(),
        };
        assert_eq!(check_availability(&server), McpStatus::Unavailable);
    }

    /// Helper: parse JSON config and return entries with command-array fixup.
    fn parse_test_config(json: &str) -> HashMap<String, McpServerEntry> {
        let val: serde_json::Value = serde_json::from_str(json).unwrap();
        let mut entries = HashMap::new();
        for key in &["mcpServers", "servers"] {
            if let Some(serde_json::Value::Object(map)) = val.get(key) {
                for (name, entry_val) in map {
                    COMMAND_EXTRA_ARGS.with(|cell| cell.borrow_mut().clear());
                    let mut entry: McpServerEntry =
                        serde_json::from_value(entry_val.clone()).unwrap();
                    entry.command_extra_args = take_command_extra_args();
                    entries.insert(name.clone(), entry);
                }
            }
        }
        entries
    }
}