llmfit 1.1.15

use llmfit_core::fit::{
    CalcConfig, FitLevel, InferenceRuntime, ModelFit, SortColumn, rank_models_by_fit_opts_col,
};
use llmfit_core::hardware::{GpuBackend, SystemSpecs};
use llmfit_core::models::{LlmModel, ModelDatabase, UseCase};
use llmfit_core::plan::{PlanRequest, estimate_model_plan_with_config};
use llmfit_core::providers::{
    DockerModelRunnerProvider, LlamaCppProvider, LmStudioProvider, MlxProvider, ModelProvider,
    OllamaProvider, RamaLamaProvider, VllmProvider,
};
use rmcp::handler::server::{router::tool::ToolRouter, wrapper::Parameters};
use rmcp::{ServerHandler, ServiceExt, tool, tool_handler, tool_router};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

use crate::serve_shared;

// --- Tool input schemas ---

#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct RecommendModelsParams {
    /// Maximum number of models to return (default: 10)
    pub limit: Option<usize>,
    /// Filter by use case: general, coding, reasoning, chat, multimodal, embedding
    pub use_case: Option<String>,
    /// Minimum fit level: perfect, good, marginal
    pub min_fit: Option<String>,
    /// Filter by runtime: mlx, llamacpp, vllm
    pub runtime: Option<String>,
    /// Filter by license string
    pub license: Option<String>,
    /// Sort by: score, tps, params, mem, ctx, date
    pub sort: Option<String>,
}

#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct SearchModelsParams {
    /// Search query (matches model name, provider, parameters, use case)
    pub query: String,
    /// Maximum number of results to return (default: 10)
    pub limit: Option<usize>,
}

#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct PlanHardwareParams {
    /// Model name to plan for
    pub model: String,
    /// Context window size (default: 8192)
    pub context: Option<u32>,
    /// Quantization level (e.g. Q4_K_M, Q8_0)
    pub quant: Option<String>,
    /// Target tokens per second
    pub target_tps: Option<f64>,
}

// --- Tool output types ---

#[derive(Serialize)]
struct RuntimeInfo {
    name: &'static str,
    installed: bool,
}

#[derive(Serialize)]
struct InstalledModel {
    name: String,
    runtime: String,
}

// --- MCP Server ---

#[derive(Clone)]
pub struct LlmfitMcpServer {
    specs: SystemSpecs,
    models: Vec<LlmModel>,
    context_limit: Option<u32>,
    /// Calculation parameters implied by `--profile`, or `None` when the
    /// server describes the machine it runs on (issue #969).
    calc_config: Option<CalcConfig>,
    node_name: String,
    tool_router: ToolRouter<Self>,
}

#[tool_handler]
impl ServerHandler for LlmfitMcpServer {}

#[tool_router]
impl LlmfitMcpServer {
    pub fn new(
        specs: SystemSpecs,
        models: Vec<LlmModel>,
        context_limit: Option<u32>,
        calc_config: Option<CalcConfig>,
        node_name: String,
    ) -> Self {
        Self {
            specs,
            models,
            context_limit,
            calc_config,
            node_name,
            tool_router: Self::tool_router(),
        }
    }

    /// Get this node's hardware specifications including RAM, GPU, CPU details
    #[tool(
        name = "get_system_specs",
        description = "Get hardware specs for this node (RAM, GPU, CPU)"
    )]
    async fn get_system_specs(&self) -> String {
        let result = serde_json::json!({
            "node": self.node_name,
            "system": serve_shared::system_json(&self.specs),
        });
        serde_json::to_string_pretty(&result).unwrap_or_default()
    }

    /// Recommend LLM models that fit this system's hardware
    #[tool(
        name = "recommend_models",
        description = "Recommend LLM models that fit this system's hardware, with optional filtering by use case, fit level, runtime, and license"
    )]
    async fn recommend_models(&self, params: Parameters<RecommendModelsParams>) -> String {
        let params = params.0;
        let limit = params.limit.unwrap_or(10);
        let sort_column = parse_sort(params.sort.as_deref());
        let min_fit = parse_min_fit(params.min_fit.as_deref());
        let runtime_filter = parse_runtime(params.runtime.as_deref());
        let use_case_filter = parse_use_case(params.use_case.as_deref());

        let mut fits = self.analyze_all();

        if let Some(min) = min_fit {
            fits.retain(|f| fit_at_least(f.fit_level, min));
        } else {
            fits.retain(|f| f.fit_level != FitLevel::TooTight);
        }

        if let Some(rt) = runtime_filter {
            fits.retain(|f| f.runtime == rt);
        }

        if let Some(uc) = use_case_filter {
            fits.retain(|f| f.use_case == uc);
        }

        if let Some(ref lic) = params.license {
            fits.retain(|f| llmfit_core::models::matches_license_filter(&f.model.license, lic));
        }

        let total = fits.len();
        let mut ranked = rank_models_by_fit_opts_col(fits, false, sort_column);
        ranked.truncate(limit);

        let result = serde_json::json!({
            "total_models": total,
            "returned_models": ranked.len(),
            "models": ranked.iter().map(serve_shared::fit_to_json).collect::<Vec<_>>(),
        });
        serde_json::to_string_pretty(&result).unwrap_or_default()
    }

    /// Search for models by name, provider, or use case
    #[tool(
        name = "search_models",
        description = "Search for LLM models by name, provider, or use case"
    )]
    async fn search_models(&self, params: Parameters<SearchModelsParams>) -> String {
        let params = params.0;
        let limit = params.limit.unwrap_or(10);
        let search_lower = params.query.to_lowercase();

        let mut fits = self.analyze_all();
        fits.retain(|f| {
            f.model.name.to_lowercase().contains(&search_lower)
                || f.model.provider.to_lowercase().contains(&search_lower)
                || f.model
                    .parameter_count
                    .to_lowercase()
                    .contains(&search_lower)
                || f.model.use_case.to_lowercase().contains(&search_lower)
                || f.use_case.label().to_lowercase().contains(&search_lower)
        });
        fits.retain(|f| f.fit_level != FitLevel::TooTight);

        let total = fits.len();
        let mut ranked = rank_models_by_fit_opts_col(fits, false, SortColumn::Score);
        ranked.truncate(limit);

        let result = serde_json::json!({
            "query": params.query,
            "total_matches": total,
            "returned_models": ranked.len(),
            "models": ranked.iter().map(serve_shared::fit_to_json).collect::<Vec<_>>(),
        });
        serde_json::to_string_pretty(&result).unwrap_or_default()
    }

    /// Plan hardware requirements for running a specific model
    #[tool(
        name = "plan_hardware",
        description = "Plan hardware requirements for running a specific model at a given context window size"
    )]
    async fn plan_hardware(&self, params: Parameters<PlanHardwareParams>) -> String {
        let params = params.0;
        let model = self
            .models
            .iter()
            .find(|m| m.name.eq_ignore_ascii_case(&params.model));

        let Some(model) = model else {
            return serde_json::json!({
                "error": format!("model '{}' not found", params.model),
            })
            .to_string();
        };

        let request = PlanRequest {
            context: params.context.unwrap_or(8192),
            quant: params.quant,
            target_tps: params.target_tps,
            kv_quant: None,
        };

        let config = self.calc_config.clone().unwrap_or_default();
        match estimate_model_plan_with_config(model, &request, &self.specs, &config) {
            Ok(estimate) => {
                serde_json::to_string_pretty(&serde_json::json!(estimate)).unwrap_or_default()
            }
            Err(e) => serde_json::json!({ "error": e }).to_string(),
        }
    }

    /// Check which inference runtimes are installed and available
    #[tool(
        name = "get_runtimes",
        description = "Check which local inference runtimes (Ollama, MLX, llama.cpp, vLLM, etc.) are installed"
    )]
    async fn get_runtimes(&self) -> String {
        let mut set = tokio::task::JoinSet::new();
        set.spawn_blocking(|| RuntimeInfo {
            name: "ollama",
            installed: OllamaProvider::new().is_available(),
        });
        set.spawn_blocking(|| RuntimeInfo {
            name: "mlx",
            installed: MlxProvider::new().is_available(),
        });
        set.spawn_blocking(|| RuntimeInfo {
            name: "llamacpp",
            installed: LlamaCppProvider::new().is_available(),
        });
        set.spawn_blocking(|| RuntimeInfo {
            name: "docker_model_runner",
            installed: DockerModelRunnerProvider::new().is_available(),
        });
        set.spawn_blocking(|| RuntimeInfo {
            name: "lmstudio",
            installed: LmStudioProvider::new().is_available(),
        });
        set.spawn_blocking(|| RuntimeInfo {
            name: "vllm",
            installed: VllmProvider::new().is_available(),
        });
        set.spawn_blocking(|| RuntimeInfo {
            name: "ramalama",
            installed: RamaLamaProvider::new().detect_with_installed().0,
        });

        let mut runtimes = Vec::new();
        while let Some(result) = set.join_next().await {
            if let Ok(info) = result {
                runtimes.push(info);
            }
        }

        let result = serde_json::json!({ "runtimes": runtimes });
        serde_json::to_string_pretty(&result).unwrap_or_default()
    }

    /// List models currently installed in local inference runtimes
    #[tool(
        name = "get_installed_models",
        description = "List models currently installed in local inference runtimes (Ollama, MLX, llama.cpp, etc.)"
    )]
    async fn get_installed_models(&self) -> String {
        let mut set = tokio::task::JoinSet::new();
        set.spawn_blocking(|| {
            let p = OllamaProvider::new();
            ("ollama", p.is_available(), p.installed_models())
        });
        set.spawn_blocking(|| {
            let p = MlxProvider::new();
            ("mlx", p.is_available(), p.installed_models())
        });
        set.spawn_blocking(|| {
            let p = LlamaCppProvider::new();
            ("llamacpp", p.is_available(), p.installed_models())
        });
        set.spawn_blocking(|| {
            let p = DockerModelRunnerProvider::new();
            (
                "docker_model_runner",
                p.is_available(),
                p.installed_models(),
            )
        });
        set.spawn_blocking(|| {
            let p = LmStudioProvider::new();
            ("lmstudio", p.is_available(), p.installed_models())
        });
        set.spawn_blocking(|| {
            let p = VllmProvider::new();
            ("vllm", p.is_available(), p.installed_models())
        });
        set.spawn_blocking(|| {
            let p = RamaLamaProvider::new();
            let (available, installed, _) = p.detect_with_installed();
            ("ramalama", available, installed)
        });

        let mut detected = Vec::new();
        while let Some(result) = set.join_next().await {
            if let Ok((name, available, installed)) = result
                && available
            {
                detected.push((name, installed));
            }
        }

        let models = installed_models_from_results(detected);
        let result = serde_json::json!({ "models": models });
        serde_json::to_string_pretty(&result).unwrap_or_default()
    }
}

impl LlmfitMcpServer {
    fn analyze_all(&self) -> Vec<ModelFit> {
        let is_apple_silicon = self.specs.backend == GpuBackend::Metal && self.specs.unified_memory;
        let mut fits: Vec<ModelFit> =
            llmfit_core::analysis::rankable_models(&self.models, &self.specs)
                .map(|m| {
                    llmfit_core::analysis::analyze_with_optional_config(
                        m,
                        &self.specs,
                        self.context_limit,
                        self.calc_config.as_ref(),
                    )
                })
                .collect();

        if !is_apple_silicon {
            fits.retain(|f| !f.model.is_mlx_only());
        }

        fits
    }
}

fn installed_models_from_results(
    results: Vec<(&'static str, std::collections::HashSet<String>)>,
) -> Vec<InstalledModel> {
    results
        .into_iter()
        .flat_map(|(runtime, models)| {
            models.into_iter().map(move |name| InstalledModel {
                name,
                runtime: runtime.to_string(),
            })
        })
        .collect()
}

// --- Entry point ---

pub fn run_mcp_server(
    overrides: &super::HardwareOverrides,
    context_limit: Option<u32>,
) -> Result<(), String> {
    let (specs, calc_config) = super::detect_specs_and_config(overrides);
    let db = ModelDatabase::new();
    let models = db.get_all_models().clone();
    let node_name = std::env::var("HOSTNAME")
        .ok()
        .filter(|v| !v.trim().is_empty())
        .unwrap_or_else(|| "unknown-node".to_string());

    let server = LlmfitMcpServer::new(specs, models, context_limit, calc_config, node_name);

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .map_err(|e| format!("failed to start tokio runtime: {e}"))?;

    runtime.block_on(async move {
        let transport = rmcp::transport::io::stdio();
        let service = server
            .serve(transport)
            .await
            .map_err(|e| format!("MCP server error: {e}"))?;
        service
            .waiting()
            .await
            .map_err(|e| format!("MCP server terminated: {e}"))?;
        Ok(())
    })
}

// --- Helper parsers ---

fn parse_sort(raw: Option<&str>) -> SortColumn {
    match raw.unwrap_or("score").trim().to_lowercase().as_str() {
        "tps" | "tokens" | "throughput" => SortColumn::Tps,
        "params" | "parameters" => SortColumn::Params,
        "mem" | "memory" | "mem_pct" => SortColumn::MemPct,
        "ctx" | "context" => SortColumn::Ctx,
        "date" | "release" => SortColumn::ReleaseDate,
        "use" | "use_case" => SortColumn::UseCase,
        _ => SortColumn::Score,
    }
}

fn parse_min_fit(raw: Option<&str>) -> Option<FitLevel> {
    raw.map(|s| match s.trim().to_lowercase().as_str() {
        "perfect" => FitLevel::Perfect,
        "good" => FitLevel::Good,
        "marginal" => FitLevel::Marginal,
        _ => FitLevel::Marginal,
    })
}

fn parse_runtime(raw: Option<&str>) -> Option<InferenceRuntime> {
    raw.and_then(|s| match s.trim().to_lowercase().as_str() {
        "mlx" => Some(InferenceRuntime::Mlx),
        "llamacpp" | "llama.cpp" | "llama_cpp" => Some(InferenceRuntime::LlamaCpp),
        "vllm" => Some(InferenceRuntime::Vllm),
        _ => None,
    })
}

fn parse_use_case(raw: Option<&str>) -> Option<UseCase> {
    raw.and_then(|s| match s.trim().to_lowercase().as_str() {
        "coding" | "code" => Some(UseCase::Coding),
        "reasoning" | "reason" => Some(UseCase::Reasoning),
        "chat" => Some(UseCase::Chat),
        "multimodal" | "vision" => Some(UseCase::Multimodal),
        "embedding" | "embed" => Some(UseCase::Embedding),
        "general" => Some(UseCase::General),
        _ => None,
    })
}

fn fit_at_least(actual: FitLevel, minimum: FitLevel) -> bool {
    let rank = |fit: FitLevel| match fit {
        FitLevel::Perfect => 3,
        FitLevel::Good => 2,
        FitLevel::Marginal => 1,
        FitLevel::TooTight => 0,
    };
    rank(actual) >= rank(minimum)
}

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

    fn test_server() -> LlmfitMcpServer {
        server_with(SystemSpecs::detect(), None)
    }

    fn server_with(specs: SystemSpecs, calc_config: Option<CalcConfig>) -> LlmfitMcpServer {
        LlmfitMcpServer::new(
            specs,
            ModelDatabase::new().get_all_models().clone(),
            None,
            calc_config,
            "test-node".to_string(),
        )
    }

    /// A fixed 128 GB unified-memory box, so estimator assertions don't depend
    /// on whatever hardware the test happens to run on.
    fn unified_specs() -> SystemSpecs {
        SystemSpecs {
            total_ram_gb: 128.0,
            available_ram_gb: 120.0,
            total_cpu_cores: 16,
            cpu_name: "Test CPU".to_string(),
            has_gpu: true,
            gpu_vram_gb: Some(128.0),
            total_gpu_vram_gb: Some(128.0),
            gpu_available_gb: None,
            gpu_name: Some("Test Unified GPU".to_string()),
            gpu_count: 1,
            unified_memory: true,
            backend: GpuBackend::Rocm,
            gpus: vec![],
            cluster_mode: false,
            cluster_node_count: 0,
        }
    }

    /// MCP clients rank from `analyze_all`, so the sanitization gate has to
    /// hold here too — an agent must not be handed a speculative-decoding
    /// draft head that the CLI hides (issue #969, problem 3).
    #[test]
    fn analyze_all_excludes_sanitization_demoted_entries() {
        let db = ModelDatabase::new();
        let demoted: std::collections::HashSet<&str> = db
            .get_all_models()
            .iter()
            .filter(|m| m.is_sanitization_demoted())
            .map(|m| m.name.as_str())
            .collect();
        assert!(
            !demoted.is_empty(),
            "catalog has no demoted entries to hide"
        );

        let fits = server_with(unified_specs(), None).analyze_all();

        assert!(!fits.is_empty(), "expected some ranked fits");
        let leaked: Vec<&str> = fits
            .iter()
            .map(|f| f.model.name.as_str())
            .filter(|name| demoted.contains(name))
            .collect();
        assert!(leaked.is_empty(), "demoted models reached MCP: {leaked:?}");
    }

    /// A profile's bandwidth has to reach the estimator, not just the reported
    /// capacity, or MCP publishes this host's speeds under the profile's name.
    #[test]
    fn analyze_all_estimates_from_the_profile_config() {
        let config = CalcConfig {
            gpu_bandwidth_gbps_override: Some(777.0),
            ..CalcConfig::default()
        };

        let bandwidths: Vec<f64> = server_with(unified_specs(), Some(config))
            .analyze_all()
            .iter()
            .filter_map(|f| f.estimate_basis.gpu_bandwidth_gbps)
            .collect();

        assert!(
            !bandwidths.is_empty() && bandwidths.iter().all(|bw| *bw == 777.0),
            "every GPU-path estimate should use the profile bandwidth, got {bandwidths:?}"
        );
    }

    #[test]
    fn runtime_discovery_includes_ramalama() {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("test runtime should build");
        let output = runtime.block_on(test_server().get_runtimes());
        let value: Value = serde_json::from_str(&output).expect("runtime output should be JSON");
        let runtimes = value
            .get("runtimes")
            .and_then(Value::as_array)
            .expect("runtime output should contain an array");

        assert!(
            runtimes
                .iter()
                .any(|runtime| { runtime.get("name").and_then(Value::as_str) == Some("ramalama") })
        );
    }

    #[test]
    fn installed_model_discovery_returns_json_when_providers_are_unavailable() {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("test runtime should build");
        let output = runtime.block_on(test_server().get_installed_models());
        let value: Value =
            serde_json::from_str(&output).expect("installed model output should be JSON");
        let models = value
            .get("models")
            .and_then(Value::as_array)
            .expect("installed model output should contain an array");

        assert!(models.iter().all(|model| {
            model.get("name").and_then(Value::as_str).is_some()
                && model.get("runtime").and_then(Value::as_str).is_some()
        }));
    }

    #[test]
    fn installed_model_discovery_preserves_ramalama_model_identity() {
        let models = installed_models_from_results(vec![(
            "ramalama",
            std::collections::HashSet::from(["model-x".to_string()]),
        )]);

        assert_eq!(models.len(), 1);
        assert_eq!(models[0].runtime, "ramalama");
        assert_eq!(models[0].name, "model-x");
    }
}