llm_manager/backend/

server.rs

use std::fmt::Display;
use std::net::IpAddr;
use std::process::Stdio;
use std::str::FromStr;
use std::sync::LazyLock;
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::Command;
use tokio::sync::mpsc;
use tracing::{info, warn};

use crate::config::Config;
use crate::models::{
    DiscoveredModel, ModelSettings, RopeScaling, ServerMetrics, clean_host, strip_gguf,
};

/// Client for health checks (short timeout).
static HEALTH_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
    reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(1))
        .build()
        .unwrap()
});

static HTTP_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(reqwest::Client::new);

/// Cached VRAM metrics from system tools (nvidia-smi / amdgpu_top).
struct VramCache {
    nvidia: Option<(u64, u64)>,
    amd: Option<(u64, u64)>,
    stale: bool,
}

static VRAM_CACHE: LazyLock<std::sync::Mutex<VramCache>> = LazyLock::new(|| {
    std::sync::Mutex::new(VramCache {
        nvidia: None,
        amd: None,
        stale: true,
    })
});

/// Invalidate the VRAM cache so the next metrics poll re-queries system tools.
pub fn invalidate_vram_cache() {
    let mut cache = VRAM_CACHE.lock().unwrap();
    cache.nvidia = None;
    cache.amd = None;
    cache.stale = true;
}

/// Manages a single llama.cpp server process.
#[derive(Clone)]
pub struct ServerHandle {
    pub port: u16,
    pub host: String,
    pub pid: u32,
    pub kill_tx: mpsc::Sender<()>,
}

/// Helper: add an argument to both the Command and the display parts list.
fn push_arg(cmd: &mut Command, parts: &mut Vec<String>, name: &str, value: impl Display) {
    let val_str = value.to_string();
    cmd.arg(name).arg(&val_str);
    parts.push(name.to_string());
    // Quote values that may contain spaces for shell safety
    if val_str.contains(' ') || val_str.contains(';') || val_str.contains('"') {
        parts.push(format!(
            "\"{}\"",
            val_str.replace('\\', "\\\\").replace('"', "\\\"")
        ));
    } else {
        parts.push(val_str);
    }
}

/// Helper: add a flag (argument without value) to both the Command and display parts.
fn push_flag(cmd: &mut Command, parts: &mut Vec<String>, name: &str) {
    cmd.arg(name);
    parts.push(name.to_string());
}

fn push_gpu_layers(cmd: &mut Command, parts: &mut Vec<String>, settings: &ModelSettings) {
    match settings.gpu_layers_mode {
        crate::models::GpuLayersMode::Specific(n) => push_arg(cmd, parts, "-ngl", n),
        crate::models::GpuLayersMode::All => push_arg(cmd, parts, "-ngl", "999"),
        crate::models::GpuLayersMode::Auto => {}
    }
}

fn push_spec_decoding(cmd: &mut Command, parts: &mut Vec<String>, settings: &ModelSettings) {
    if !settings.spec_type.is_empty() {
        push_arg(cmd, parts, "--spec-type", &settings.spec_type);
        if settings.draft_tokens > 0 {
            push_arg(cmd, parts, "--spec-draft-n-max", settings.draft_tokens);
        }
    }
}

/// Build the full llama-server command line from settings.
/// Returns (Command, display_string) where the string is suitable for logging.
pub fn build_server_cmd(
    binary: &std::path::Path,
    model: Option<&DiscoveredModel>,
    settings: &ModelSettings,
    config: &Config,
    server_mode: crate::models::ServerMode,
    router_max_models: u32,
    model_config_used: bool,
) -> (Command, String) {
    let mut cmd = Command::new(binary);
    let mut parts: Vec<String> = vec![binary.display().to_string()];

    // ── Model ───────────────────────────────────────────────
    match server_mode {
        crate::models::ServerMode::Normal => {
            if let Some(model) = model {
                push_arg(&mut cmd, &mut parts, "-m", model.path.display());
                // Add alias for router mode identification (uses the unique relative path)
                push_arg(&mut cmd, &mut parts, "--alias", &model.display_name);
            }
        }
        crate::models::ServerMode::Router => {
            // Router mode: no model in CLI, use /load API to load models
            if router_max_models > 0 {
                push_arg(&mut cmd, &mut parts, "--models-max", router_max_models);
            }
            // Always pass --models-dir in router mode (global config setting)
            if let Some(dir) = config.models_dirs.first() {
                push_arg(&mut cmd, &mut parts, "--models-dir", dir.display());
            }
        }
        crate::models::ServerMode::Bench => {
            // Should not be reached as Bench uses build_bench_cmd
        }
        crate::models::ServerMode::BenchTune => {
            // Should not be reached as BenchTune uses benchmark tuning function
        }
    }

    // Parse GGUF metadata for arch-specific override-kv keys
    let gguf_meta = model
        .map(|m| crate::models::GgufMetadata::from_path(&m.path))
        .transpose();

    // ── Loading ──────────────────────────────────────────────
    push_arg(&mut cmd, &mut parts, "--threads", settings.threads);
    push_arg(
        &mut cmd,
        &mut parts,
        "--threads-batch",
        settings.threads_batch,
    );
    let effective_ctx = (settings.context_length as f64 * settings.rope_scale as f64) as u32;
    if !model_config_used {
        push_arg(&mut cmd, &mut parts, "--ctx-size", effective_ctx);
    }
    push_arg(&mut cmd, &mut parts, "--ubatch-size", settings.ubatch_size);
    if let Some(n) = settings.max_concurrent_predictions {
        push_arg(&mut cmd, &mut parts, "--parallel", n);
    }

    push_flag(&mut cmd, &mut parts, "--no-warmup");

    push_spec_decoding(&mut cmd, &mut parts, settings);

    if let Some(cache_k) = settings.cache_type_k {
        push_arg(&mut cmd, &mut parts, "--cache-type-k", cache_k);
    }
    if let Some(cache_v) = settings.cache_type_v {
        push_arg(&mut cmd, &mut parts, "--cache-type-v", cache_v);
    }

    if settings.keep != 0 {
        push_arg(&mut cmd, &mut parts, "--keep", settings.keep);
    }
    if settings.swa_full {
        push_flag(&mut cmd, &mut parts, "--swa-full");
    }
    if settings.mlock {
        push_flag(&mut cmd, &mut parts, "--mlock");
    }
    if !settings.mmap {
        push_flag(&mut cmd, &mut parts, "--no-mmap");
    }
    if settings.numa != Default::default() {
        push_arg(&mut cmd, &mut parts, "--numa", settings.numa.to_string());
    }
    if settings.kv_cache_offload {
        push_flag(&mut cmd, &mut parts, "--kv-offload");
    }

    // ── GPU ──────────────────────────────────────────────────
    push_gpu_layers(&mut cmd, &mut parts, settings);

    if settings.split_mode != Default::default() {
        push_arg(
            &mut cmd,
            &mut parts,
            "--split-mode",
            settings.split_mode.to_string(),
        );
    }
    if !settings.tensor_split.is_empty() {
        push_arg(
            &mut cmd,
            &mut parts,
            "--tensor-split",
            &settings.tensor_split,
        );
    }
    if settings.main_gpu != 0 {
        push_arg(&mut cmd, &mut parts, "--main-gpu", settings.main_gpu);
    }
    if settings.fit {
        push_arg(&mut cmd, &mut parts, "--fit", "on");
    }

    if let Some(ref lora) = settings.lora {
        push_arg(&mut cmd, &mut parts, "--lora", lora.display());
    }
    if let Some((ref lora, scale)) = settings.lora_scaled {
        push_arg(
            &mut cmd,
            &mut parts,
            "--lora-scaled",
            format!("{}:{}", lora.display(), scale),
        );
    }

    let mut rpc_list = Vec::new();
    if !settings.rpc.is_empty() {
        rpc_list.push(settings.rpc.clone());
    }
    for worker in &config.rpc_workers {
        if worker.selected {
            if IpAddr::from_str(&worker.ip).is_ok() {
                rpc_list.push(format!("{}:{}", worker.ip, worker.port));
            }
        }
    }

    if !rpc_list.is_empty() {
        let joined_rpc = rpc_list.join(",");
        push_arg(&mut cmd, &mut parts, "--rpc", joined_rpc);
    }

    if settings.embedding {
        push_flag(&mut cmd, &mut parts, "--embedding");
    }

    if settings.expert_count > 0 {
        let arch = gguf_meta
            .as_ref()
            .ok()
            .and_then(|opt| opt.as_ref())
            .map(|m| m.arch.as_str())
            .unwrap_or("llama");
        let arch = arch
            .chars()
            .filter(|c| c.is_ascii_alphanumeric() || *c == '_')
            .collect::<String>();
        push_arg(
            &mut cmd,
            &mut parts,
            "--override-kv",
            format!(
                "{}.expert_used_count=int:int:{}",
                arch, settings.expert_count
            ),
        );
    }

    push_arg(
        &mut cmd,
        &mut parts,
        "-fa",
        if settings.flash_attn { "on" } else { "off" },
    );

    if settings.jinja {
        push_flag(&mut cmd, &mut parts, "--jinja");
    }

    // Apply chat template: file path, built-in name, or auto-detect from arch
    if let Some(ref t) = settings.chat_template {
        if t.ends_with(".jinja") {
            push_arg(&mut cmd, &mut parts, "--chat-template-file", t);
        } else {
            push_arg(&mut cmd, &mut parts, "--chat-template", t);
        }
    } else if settings.auto_chat_template {
        let arch = gguf_meta
            .as_ref()
            .ok()
            .and_then(|opt| opt.as_ref())
            .map(|m| m.arch.as_str())
            .unwrap_or("llama");
        if let Some(template) = crate::models::arch_to_chat_template(arch) {
            push_arg(&mut cmd, &mut parts, "--chat-template", template);
        }
    }

    // Inject system prompt via chat template kwargs when it is not empty
    if !settings.system_prompt.is_empty() {
        let mut merged = serde_json::Map::new();
        if let Some(ref kwargs) = settings.chat_template_kwargs
            && let Ok(obj) = serde_json::from_str::<serde_json::Value>(kwargs)
            && let serde_json::Value::Object(map) = obj
        {
            for (k, v) in map {
                merged.insert(k, v);
            }
        }
        // Warn if user's kwargs already defines system_prompt
        if merged.contains_key("system_prompt") {
            tracing::warn!(
                "chat_template_kwargs already contains 'system_prompt', will be overridden by settings"
            );
        }
        merged.insert(
            "system_prompt".to_string(),
            serde_json::Value::String(settings.system_prompt.clone()),
        );
        push_arg(
            &mut cmd,
            &mut parts,
            "--chat-template-kwargs",
            serde_json::to_string(&merged).unwrap(),
        );
    } else if let Some(ref kwargs) = settings.chat_template_kwargs {
        push_arg(&mut cmd, &mut parts, "--chat-template-kwargs", kwargs);
    }

    // ── Sampling ─────────────────────────────────────────────
    if settings.seed != -1 {
        push_arg(&mut cmd, &mut parts, "--seed", settings.seed);
    }
    if let Some(max_tokens) = settings.max_tokens {
        push_arg(&mut cmd, &mut parts, "--n-predict", max_tokens);
    }
    push_arg(
        &mut cmd,
        &mut parts,
        "--temp",
        format!("{:.2}", settings.temperature),
    );

    push_arg(&mut cmd, &mut parts, "--top-k", settings.top_k);

    push_arg(
        &mut cmd,
        &mut parts,
        "--top-p",
        format!("{:.2}", settings.top_p),
    );

    push_arg(
        &mut cmd,
        &mut parts,
        "--min-p",
        format!("{:.2}", settings.min_p),
    );

    push_arg(
        &mut cmd,
        &mut parts,
        "--typical",
        format!("{:.2}", settings.typical_p),
    );

    if settings.mirostat != Default::default() {
        push_arg(
            &mut cmd,
            &mut parts,
            "--mirostat",
            settings.mirostat.to_string(),
        );
        push_arg(
            &mut cmd,
            &mut parts,
            "--mirostat-lr",
            format!("{:.2}", settings.mirostat_lr),
        );
        push_arg(
            &mut cmd,
            &mut parts,
            "--mirostat-ent",
            format!("{:.2}", settings.mirostat_ent),
        );
    }

    if settings.ignore_eos {
        push_flag(&mut cmd, &mut parts, "--ignore-eos");
    }

    if !settings.samplers.0.is_empty() {
        push_arg(
            &mut cmd,
            &mut parts,
            "--samplers",
            settings.samplers.to_string(),
        );
    }

    if let Some(frequency) = settings.frequency_penalty {
        push_arg(
            &mut cmd,
            &mut parts,
            "--frequency-penalty",
            format!("{:.2}", frequency),
        );
    }

    if settings.dry_multiplier != 0.0 {
        push_arg(
            &mut cmd,
            &mut parts,
            "--dry-multiplier",
            format!("{:.2}", settings.dry_multiplier),
        );
        push_arg(
            &mut cmd,
            &mut parts,
            "--dry-base",
            format!("{:.2}", settings.dry_base),
        );
        push_arg(
            &mut cmd,
            &mut parts,
            "--dry-allowed-length",
            settings.dry_allowed_length,
        );
        push_arg(
            &mut cmd,
            &mut parts,
            "--dry-penalty-last-n",
            settings.dry_penalty_last_n,
        );
    }

    // ── RoPE ─────────────────────────────────────────────────
    let rope_scaling = if settings.rope_yarn_enabled {
        RopeScaling::Yarn
    } else {
        settings.rope_scaling
    };
    if rope_scaling != Default::default() {
        push_arg(
            &mut cmd,
            &mut parts,
            "--rope-scaling",
            rope_scaling.to_string(),
        );
    }
    if settings.rope_scale != 1.0 {
        push_arg(
            &mut cmd,
            &mut parts,
            "--rope-scale",
            format!("{:.2}", settings.rope_scale),
        );
    }
    if settings.rope_freq_base != 0.0 {
        push_arg(
            &mut cmd,
            &mut parts,
            "--rope-freq-base",
            format!("{:.2}", settings.rope_freq_base),
        );
    }
    if settings.rope_freq_scale != 1.0 {
        push_arg(
            &mut cmd,
            &mut parts,
            "--rope-freq-scale",
            format!("{:.2}", settings.rope_freq_scale),
        );
    }

    if settings.rope_yarn_enabled
        && settings.rope_scale > 1.0
        && let Some(meta) = gguf_meta.as_ref().ok().and_then(|x| x.as_ref())
    {
        push_arg(
            &mut cmd,
            &mut parts,
            "--override-kv",
            format!("{}.context_length=int:{}", meta.arch, effective_ctx),
        );
        let orig_ctx = meta.n_ctx_train;
        push_arg(&mut cmd, &mut parts, "--yarn-orig-ctx", orig_ctx);
    }

    let resolved_host = clean_host(&settings.host);
    push_arg(&mut cmd, &mut parts, "--host", resolved_host);
    push_arg(&mut cmd, &mut parts, "--port", settings.port);
    push_arg(&mut cmd, &mut parts, "--timeout", settings.timeout);

    push_flag(&mut cmd, &mut parts, "--metrics");
    if !settings.cache_prompt {
        push_flag(&mut cmd, &mut parts, "--no-cache-prompt");
    }
    if settings.cache_reuse != 0 {
        push_arg(&mut cmd, &mut parts, "--cache-reuse", settings.cache_reuse);
    }
    if !settings.webui {
        push_flag(&mut cmd, &mut parts, "--no-webui");
    }
    let lv = match config.default.log_level.as_str() {
        "error" => 0,
        "warn" => 1,
        "info" => 2,
        "debug" => 3,
        "trace" => 4,
        _ => 1,
    };
    push_arg(&mut cmd, &mut parts, "-lv", lv);

    // ── General ──────────────────────────────────────────────

    let display = parts.join(" ");
    (cmd, display)
}

/// Build the full llama-bench command line.
pub fn build_bench_cmd(
    binary: &std::path::Path,
    model: &DiscoveredModel,
    settings: &ModelSettings,
) -> (Command, String) {
    let mut cmd = Command::new(binary);
    let mut parts: Vec<String> = vec![binary.display().to_string()];

    push_arg(&mut cmd, &mut parts, "-m", model.path.display());
    push_arg(&mut cmd, &mut parts, "-t", settings.threads);
    push_arg(&mut cmd, &mut parts, "-b", settings.batch_size);

    push_gpu_layers(&mut cmd, &mut parts, settings);

    if settings.flash_attn {
        push_arg(&mut cmd, &mut parts, "-fa", "1");
    }

    push_flag(&mut cmd, &mut parts, "--progress");

    let display = parts.join(" ");
    (cmd, display)
}

/// Spawn a llama.cpp server process (single model or router).
/// Returns (ServerHandle, command_string) where command_string is the full CLI.
pub struct SpawnServerRequest<'a> {
    pub config: &'a Config,
    pub model: Option<&'a DiscoveredModel>,
    pub settings: &'a ModelSettings,
    pub log_tx: mpsc::Sender<String>,
    pub progress_tx: Option<tokio::sync::broadcast::Sender<crate::models::DownloadState>>,
    pub server_mode: crate::models::ServerMode,
    pub router_max_models: u32,
    pub exit_tx: mpsc::Sender<()>,
}

pub async fn spawn_server(req: SpawnServerRequest<'_>) -> Result<(ServerHandle, String), String> {
    let SpawnServerRequest {
        config,
        model,
        settings,
        log_tx,
        progress_tx,
        server_mode,
        router_max_models,
        exit_tx,
    } = req;
    if server_mode != crate::models::ServerMode::Bench
        && server_mode != crate::models::ServerMode::BenchTune
    {
        let port = settings.port;
        // Check if port is already in use (bind to the same host the server will use)
        let resolved_host = clean_host(&settings.host);
        if std::net::TcpListener::bind(format!("{}:{}", resolved_host, port)).is_err() {
            return Err(format!("Port {} is already in use", port));
        }
    }

    // BenchTune mode is handled separately in app.process_pending_spawn()
    // and should never reach this function.
    if server_mode == crate::models::ServerMode::BenchTune {
        return Err("BenchTune mode is not supported in spawn_server".to_string());
    }

    // Resolve the backend binary (downloads if needed)
    let backend_name = if server_mode == crate::models::ServerMode::Bench {
        "llama-bench"
    } else {
        "llama-server"
    };
    let version_display = settings.get_active_backend_version_display();
    info!(
        "spawn_server: backend={}, requested_version={:?}, version_display={}",
        settings.backend,
        settings.get_active_backend_version(),
        version_display
    );
    log_tx
        .send(format!(
            "Resolving {} (v{}) binary...",
            backend_name, version_display
        ))
        .await
        .ok();
    let version_param = settings.get_active_backend_version().map(|s| s.as_str());

    let server_binary = match crate::backend::hub::resolve_backend_binary(
        settings.backend,
        version_param,
        Some(log_tx.clone()),
        progress_tx,
    )
    .await
    {
        Ok(path) => {
            info!("spawn_server: resolved binary path={}", path.display());
            path
        }
        Err(e) => {
            return Err(format!("Failed to resolve backend binary: {}", e));
        }
    };

    let binary = if server_mode == crate::models::ServerMode::Bench {
        server_binary.parent().unwrap().join("llama-bench")
    } else {
        server_binary
    };

    if !binary.exists() {
        return Err(format!("Binary not found at: {}", binary.display()));
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        if let Ok(metadata) = binary.metadata()
            && metadata.permissions().mode() & 0o111 == 0
        {
            return Err(format!("Binary is not executable: {}", binary.display()));
        }
    }

    let (mut cmd, cmd_string) = if server_mode == crate::models::ServerMode::Bench {
        if let Some(m) = model {
            build_bench_cmd(&binary, m, settings)
        } else {
            return Err("Model required for benchmark".to_string());
        }
    } else {
        build_server_cmd(
            &binary,
            model,
            settings,
            config,
            server_mode,
            router_max_models,
            false,
        )
    };

    cmd.stdout(Stdio::piped()).stderr(Stdio::piped());

    // Set platform-specific env vars so the binary can find its shared libraries
    let bin_dir = binary.parent().unwrap();
    match std::env::consts::OS {
        "windows" => {
            // On Windows, add bin_dir to PATH so llama-server.exe finds libllama.dll
            if let Ok(current) = std::env::var("PATH") {
                cmd.env("PATH", format!("{};{}", bin_dir.display(), current));
            } else {
                cmd.env("PATH", bin_dir);
            }
        }
        "macos" => {
            // On macOS, set DYLD_LIBRARY_PATH for dylib loading
            if let Ok(current) = std::env::var("DYLD_LIBRARY_PATH") {
                cmd.env(
                    "DYLD_LIBRARY_PATH",
                    format!("{}:{}", bin_dir.display(), current),
                );
            } else {
                cmd.env("DYLD_LIBRARY_PATH", bin_dir);
            }
        }
        _ => {
            // On Linux, set LD_LIBRARY_PATH for so loading
            if let Ok(current) = std::env::var("LD_LIBRARY_PATH") {
                cmd.env(
                    "LD_LIBRARY_PATH",
                    format!("{}:{}", bin_dir.display(), current),
                );
            } else {
                cmd.env("LD_LIBRARY_PATH", bin_dir);
            }
        }
    }

    info!("Spawning: {}", cmd_string);
    let _ = log_tx
        .send(format!("{}: {}", backend_name, cmd_string))
        .await;
    let mut child = cmd
        .spawn()
        .map_err(|e| format!("Failed to spawn process: {}", e))?;
    let pid = child.id().unwrap_or(0);

    let (kill_tx, mut kill_rx) = mpsc::channel(1);

    // Background task: read stdout and stderr concurrently via separate tasks.
    // Each stream gets its own task + mpsc channel so neither can block the other.
    let log_tx_inner = log_tx.clone();
    let exit_tx_inner = exit_tx.clone();
    let backend_name_upper = backend_name.to_uppercase();
    tokio::spawn(async move {
        let stdout = child.stdout.take().unwrap();
        let stderr = child.stderr.take().unwrap();

        let (stdout_tx, mut stdout_rx) = mpsc::channel::<String>(64);
        let (stderr_tx, mut stderr_rx) = mpsc::channel::<String>(64);

        // Spawn a reader task for each stream
        let mut std_out = Some(tokio::spawn(async move {
            let reader = BufReader::new(stdout).lines();
            tokio::pin!(reader);
            while let Ok(Some(line)) = reader.next_line().await {
                if stdout_tx.send(line).await.is_err() {
                    break;
                }
            }
        }));

        let mut std_err = Some(tokio::spawn(async move {
            let reader = BufReader::new(stderr).lines();
            tokio::pin!(reader);
            while let Ok(Some(line)) = reader.next_line().await {
                if stderr_tx.send(line).await.is_err() {
                    break;
                }
            }
        }));

        // Merge loop: block on whichever channel has data.
        // When both are empty, select! sleeps with zero CPU cost.
        loop {
            tokio::select! {
                _ = kill_rx.recv() => {
                    let _ = child.kill().await;
                    if let Some(h) = std_out.take() { let _ = h.await; }
                    if let Some(h) = std_err.take() { let _ = h.await; }
                    break;
                }
                line = stdout_rx.recv() => {
                    if let Some(line) = line { let _ = log_tx_inner.send(line).await; } else { break; }
                }
                line = stderr_rx.recv() => {
                    if let Some(line) = line { let _ = log_tx_inner.send(line).await; } else { break; }
                }
                else => break,
            }
        }

        // Wait for reader tasks to finish
        if let Some(h) = std_out.take() {
            let _ = h.await;
        }
        if let Some(h) = std_err.take() {
            let _ = h.await;
        }

        let exit_code = child.wait().await.ok().and_then(|s| s.code());
        let _ = exit_tx_inner.send(()).await;
        let _ = log_tx_inner
            .send(format!(
                "{} exited with code {:?}",
                backend_name_upper, exit_code
            ))
            .await;
    });

    Ok((
        ServerHandle {
            port: if server_mode == crate::models::ServerMode::Bench {
                0
            } else {
                settings.port
            },
            host: settings.host.clone(),
            pid,
            kill_tx,
        },
        cmd_string,
    ))
}

/// Check if the server is healthy and responsive.
pub async fn check_health(host: &str, port: u16) -> bool {
    let host = clean_host(host);
    let url = format!("http://{}:{}/health", host, port);

    match HEALTH_CLIENT.get(&url).send().await {
        Ok(resp) => resp.status().is_success(),
        Err(_) => false,
    }
}

/// Kill a running server.
pub async fn kill_server(handle: ServerHandle) -> Result<(), String> {
    handle
        .kill_tx
        .send(())
        .await
        .map_err(|_| "Server already stopped".to_string())
}

/// Poll metrics from the server.
pub async fn get_metrics(
    host: &str,
    port: u16,
    model_name: Option<&str>,
    pid: Option<u32>,
) -> Result<ServerMetrics, String> {
    let host = clean_host(host);
    // We prefer the /metrics endpoint as it's more stable for system info.
    // In router mode, we can specify the model via query parameter.
    let mut url = if let Some(model) = model_name {
        let name = strip_gguf(model);
        format!(
            "http://{}:{}/metrics?model={}",
            host,
            port,
            urlencoding::encode(&name)
        )
    } else {
        format!("http://{}:{}/metrics", host, port)
    };

    let mut resp = HTTP_CLIENT
        .get(&url)
        .send()
        .await
        .map_err(|e| format!("Failed to get metrics: {}", e))?;

    // If model-specific metrics fail with 404 or 400, try plain /metrics
    if (resp.status() == reqwest::StatusCode::NOT_FOUND
        || resp.status() == reqwest::StatusCode::BAD_REQUEST)
        && model_name.is_some()
    {
        url = format!("http://{}:{}/metrics", host, port);
        resp = HTTP_CLIENT
            .get(&url)
            .send()
            .await
            .map_err(|e| format!("Failed to get metrics: {}", e))?;
    }

    if !resp.status().is_success() {
        return Err(format!("Server returned {}", resp.status()));
    }

    let text = resp
        .text()
        .await
        .map_err(|e| format!("Failed to read metrics: {}", e))?;

    let mut m = ServerMetrics::default();

    let mut ctx_max_slots = 0u32;
    let mut ctx_used_slots = 0u32;
    let mut ctx_used_global = 0u32;
    let mut ctx_max_global = 0u32;

    let mut vram_used_slots = 0u64;
    let mut vram_total_slots = 0u64;
    let mut vram_used_global = 0u64;
    let mut vram_total_global = 0u64;

    for line in text.lines() {
        if line.starts_with('#') || line.is_empty() {
            continue;
        }

        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() < 2 {
            continue;
        }

        let name_with_labels = parts[0];
        let mut val = 0.0;
        for part in parts.iter().skip(1) {
            if let Ok(v) = part.parse::<f64>() {
                val = v;
                break;
            }
        }

        let is_slot = name_with_labels.contains("slot=\"") || name_with_labels.contains("pool=\"");
        let name = name_with_labels
            .split('{')
            .next()
            .unwrap_or(name_with_labels);

        match name {
            "llama_kv_cache_usage_bytes"
            | "kv_cache_usage_bytes"
            | "llama_server_kv_cache_usage_bytes"
            | "llama_server_kv_cache_used_bytes"
            | "llama_server_vram_used_bytes" => {
                if is_slot {
                    vram_used_slots += val as u64;
                } else {
                    vram_used_global = vram_used_global.max(val as u64);
                }
            }
            "llama_kv_cache_total_bytes"
            | "kv_cache_total_bytes"
            | "llama_server_kv_cache_total_bytes"
            | "llama_server_vram_total_bytes" => {
                if is_slot {
                    vram_total_slots += val as u64;
                } else {
                    vram_total_global = vram_total_global.max(val as u64);
                }
            }
            "llama_model_memory_usage_bytes"
            | "model_memory_usage_bytes"
            | "llama_server_model_memory_usage_bytes"
            | "llama_server_memory_usage_bytes"
            | "llama_server_ram_usage_bytes"
            | "llama_server_mem_used_bytes" => {
                m.ram_used = m.ram_used.max(val as u64);
            }
            "llama_kv_cache_tokens_used"
            | "kv_cache_usage_tokens"
            | "kv_cache_tokens_used"
            | "llama_server_kv_cache_tokens_used"
            | "llamacpp:n_tokens_used"
            | "llama_server_n_tokens_used"
            | "llama_server_n_past"
            | "llamacpp:n_past" => {
                if is_slot {
                    ctx_used_slots += val as u32;
                } else {
                    ctx_used_global = ctx_used_global.max(val as u32);
                }
            }
            "llama_kv_cache_tokens_total"
            | "kv_cache_total_tokens"
            | "kv_cache_tokens_total"
            | "llama_server_kv_cache_tokens_total"
            | "llamacpp:n_ctx"
            | "llamacpp:n_tokens_max"
            | "llama_server_n_ctx"
            | "llama_server_n_tokens_max" => {
                if is_slot {
                    ctx_max_slots += val as u32;
                } else {
                    ctx_max_global = ctx_max_global.max(val as u32);
                }
            }
            "llama_server_cpu_usage_percentage"
            | "cpu_usage_percentage"
            | "llama_server_cpu_usage"
            | "llama_server_cpu_percent" => {
                m.cpu_usage = m.cpu_usage.max(val);
            }
            "llamacpp:predicted_tokens_seconds"
            | "llama_server_predicted_tokens_seconds"
            | "llama_server_tps" => {
                m.tps += val;
            }
            "llamacpp:prompt_tokens_seconds"
            | "llama_server_prompt_tokens_seconds"
            | "llama_server_prompt_tps" => {
                m.prompt_tps += val;
            }
            "llamacpp:kv_cache_usage_ratio" | "llama_server_kv_cache_usage_ratio" => {
                if !is_slot && ctx_max_global > 0 {
                    ctx_used_global = ctx_used_global.max((val * ctx_max_global as f64) as u32);
                }
            }
            _ => {
                tracing::debug!("Unknown metric: {}", name);
            }
        }
    }

    // Prefer global metrics (includes model weights + KV cache) over slot-only (KV cache subset).
    m.gpu_mem_used = if vram_used_global > 0 {
        vram_used_global
    } else if vram_used_slots > 0 {
        vram_used_slots
    } else {
        0
    };
    m.gpu_mem_total = if vram_total_global > 0 {
        vram_total_global
    } else if vram_total_slots > 0 {
        vram_total_slots
    } else {
        0
    };

    // ctx_used = tokens currently in the KV cache.
    // ctx_max = the total context window size allocated by the server.
    m.ctx_used = if ctx_used_slots > 0 {
        ctx_used_slots
    } else {
        ctx_used_global
    };
    m.ctx_max = if ctx_max_slots > 0 {
        ctx_max_slots
    } else {
        ctx_max_global
    };
    // ctx_max may be overridden in poll_metrics() by the user-configured value.

    // Prefer actual GPU memory usage from nvidia-smi or amdgpu_top.
    // llama-server's kv_cache_usage_bytes only reports KV cache (typically 10%
    // of total VRAM); model weights are loaded into GPU memory but not tracked
    // by the server, so we use system-level tools to report what users see on GPUs.
    if model_name.is_none() {
        // Prefer system-level VRAM over llama-server's KV-only value.
        // System tools report actual GPU memory including model weights,
        // which is what users see on their GPUs and expect to read.
        let set_if_better = |out: &mut ServerMetrics, used: u64, total: u64| {
            if out.gpu_mem_used == 0 || used > out.gpu_mem_used {
                out.gpu_mem_used = used;
                out.gpu_mem_total = total;
            }
        };

        let cached_nvidia;
        let cached_amd;
        {
            let cache = VRAM_CACHE.lock().unwrap();
            cached_nvidia = cache.nvidia;
            cached_amd = cache.amd;
        }

        if cached_nvidia.is_some() {
            if let Some((nv_used, nv_total)) = cached_nvidia {
                set_if_better(&mut m, nv_used, nv_total);
            }
        } else {
            let nv = tokio::task::spawn_blocking(get_nvidia_vram_metrics)
                .await
                .unwrap_or(Err("spawn join error".to_string()));
            if let Ok((used, total)) = nv {
                {
                    let mut cache = VRAM_CACHE.lock().unwrap();
                    cache.nvidia = Some((used, total));
                    cache.stale = false;
                }
                set_if_better(&mut m, used, total);
            }
        }

        if m.gpu_mem_total == 0 {
            if cached_amd.is_some() {
                if let Some((amd_used, amd_total)) = cached_amd {
                    set_if_better(&mut m, amd_used, amd_total);
                }
            } else {
                // AMD fallback when nvidia-smi is not available.
                let amd = tokio::task::spawn_blocking(get_amdgpu_vram_metrics)
                    .await
                    .unwrap_or(Err("spawn join error".to_string()));
                if let Ok((used, total)) = amd {
                    {
                        let mut cache = VRAM_CACHE.lock().unwrap();
                        cache.amd = Some((used, total));
                        cache.stale = false;
                    }
                    set_if_better(&mut m, used, total);
                }
            }
        }
    } else if m.gpu_mem_used == 0 {
        // KV-only queries: use system tools as a last resort.
        let cached_nvidia;
        let cached_amd;
        {
            let cache = VRAM_CACHE.lock().unwrap();
            cached_nvidia = cache.nvidia;
            cached_amd = cache.amd;
        }

        if cached_nvidia.is_some() {
            if let Some((used, total)) = cached_nvidia {
                m.gpu_mem_used = used;
                m.gpu_mem_total = total;
            }
        } else {
            let nv = tokio::task::spawn_blocking(get_nvidia_vram_metrics)
                .await
                .unwrap_or(Err("spawn join error".to_string()));
            if let Ok((used, total)) = nv {
                {
                    let mut cache = VRAM_CACHE.lock().unwrap();
                    cache.nvidia = Some((used, total));
                    cache.stale = false;
                }
                m.gpu_mem_used = used;
                m.gpu_mem_total = total;
            }
        }

        if m.gpu_mem_used == 0 {
            if cached_amd.is_some() {
                if let Some((used, total)) = cached_amd {
                    m.gpu_mem_used = used;
                    m.gpu_mem_total = total;
                }
            } else {
                let amd = tokio::task::spawn_blocking(get_amdgpu_vram_metrics)
                    .await
                    .unwrap_or(Err("spawn join error".to_string()));
                if let Ok((used, total)) = amd {
                    {
                        let mut cache = VRAM_CACHE.lock().unwrap();
                        cache.amd = Some((used, total));
                        cache.stale = false;
                    }
                    m.gpu_mem_used = used;
                    m.gpu_mem_total = total;
                }
            }
        }
    }

    // Fallback for RAM and CPU using sysinfo (cross-platform)
    if let Some(p) = pid
        && let Ok((ram, cpu)) = get_process_metrics(p)
    {
        if m.ram_used == 0 {
            m.ram_used = ram;
        }
        if m.cpu_usage == 0.0 {
            m.cpu_usage = cpu;
        }
    }

    m.loaded = ctx_max_global > 0 || vram_used_global > 0;

    Ok(m)
}

/// Get VRAM usage using nvidia-smi
fn get_nvidia_vram_metrics() -> Result<(u64, u64), String> {
    let output = std::process::Command::new("nvidia-smi")
        .args([
            "--query-gpu=memory.used,memory.total",
            "--format=csv,noheader,nounits",
        ])
        .output()
        .map_err(|e| e.to_string())?;

    if !output.status.success() {
        return Err("nvidia-smi failed".to_string());
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let mut total_used: u64 = 0;
    let mut total_total: u64 = 0;
    for line in stdout.lines() {
        let parts: Vec<&str> = line.split(',').collect();
        if parts.len() >= 2 {
            let used = match parts[0].trim().parse::<u64>() {
                Ok(v) => v,
                Err(_) => {
                    warn!(
                        "nvidia-smi: failed to parse used memory from '{}'",
                        parts[0]
                    );
                    continue;
                }
            } * 1024
                * 1024;
            let total = match parts[1].trim().parse::<u64>() {
                Ok(v) => v,
                Err(_) => {
                    warn!(
                        "nvidia-smi: failed to parse total memory from '{}'",
                        parts[1]
                    );
                    continue;
                }
            } * 1024
                * 1024;
            total_used += used;
            total_total += total;
        }
    }
    if total_total > 0 {
        return Ok((total_used, total_total));
    }

    Err("Invalid output from nvidia-smi".to_string())
}

/// Get VRAM usage using amdgpu_top
fn get_amdgpu_vram_metrics() -> Result<(u64, u64), String> {
    let output = std::process::Command::new("amdgpu_top")
        .args(["-d", "--json"])
        .output()
        .map_err(|e| e.to_string())?;

    if !output.status.success() {
        return Err("amdgpu_top failed".to_string());
    }

    let json: serde_json::Value =
        serde_json::from_slice(&output.stdout).map_err(|e| e.to_string())?;

    // amdgpu_top --json output has a "devices" array (or sometimes just a list of objects depending on version)
    let devices = if json.is_array() {
        json.as_array()
    } else {
        json.get("devices").and_then(|d| d.as_array())
    };

    if let Some(devices) = devices
        && let Some(device) = devices.first()
    {
        // Priority 1: Check root keys (newer amdgpu_top format as provided by user)
        // "VRAM Usage Size": 3070128128, "VRAM Size": 8589934592
        let root_used = device.get("VRAM Usage Size").and_then(|v| v.as_u64());
        let root_total = device.get("VRAM Size").and_then(|v| v.as_u64());

        if let (Some(used), Some(total)) = (root_used, root_total)
            && total > 0
        {
            return Ok((used, total));
        }

        // Priority 2: Check nested VRAM object (alternative format)
        let vram_obj = device.get("VRAM");
        if let Some(vram) = vram_obj {
            // Check if it's the "Total VRAM Usage" format (usually MiB)
            let nested_used = vram
                .get("Total VRAM Usage")
                .and_then(|v| v.get("value").or(Some(v)))
                .and_then(|v| v.as_u64());
            let nested_total = vram
                .get("Total VRAM")
                .and_then(|v| v.get("value").or(Some(v)))
                .and_then(|v| v.as_u64());

            if let (Some(used), Some(total)) = (nested_used, nested_total) {
                // These are usually in MiB if they have a "unit" field
                let multiplier = if vram.get("Total VRAM").and_then(|v| v.get("unit")).is_some() {
                    1024 * 1024
                } else {
                    1
                };
                if total > 0 {
                    return Ok((used * multiplier, total * multiplier));
                }
            }
        }

        // Priority 3: Check vram_usage key (older format)
        let vram_usage = device.get("vram_usage");
        if let Some(vram) = vram_usage {
            let used = vram
                .get("VRAM")
                .or_else(|| vram.get("usage"))
                .and_then(|v| v.get("value").or(Some(v)))
                .and_then(|v| v.as_u64())
                .unwrap_or(0);
            let total = vram
                .get("TotalVRAM")
                .or_else(|| vram.get("total"))
                .and_then(|v| v.get("value").or(Some(v)))
                .and_then(|v| v.as_u64())
                .unwrap_or(0);

            if total > 0 {
                return Ok((used * 1024 * 1024, total * 1024 * 1024));
            }
        }
    }

    Err("Could not find VRAM info in amdgpu_top output".to_string())
}

/// Cross-platform: Get RAM (RSS) and CPU usage for a PID.
/// Uses a persistent System instance so sysinfo can compute accurate
/// CPU deltas across calls (first call on a fresh instance is always 0).
fn get_process_metrics(pid: u32) -> Result<(u64, f64), String> {
    use std::sync::{LazyLock, Mutex};
    use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System};

    static SYS: LazyLock<Mutex<System>> = LazyLock::new(|| {
        Mutex::new(System::new_with_specifics(
            RefreshKind::everything()
                .with_processes(ProcessRefreshKind::nothing().with_cpu().with_memory()),
        ))
    });

    let mut sys = SYS.lock().unwrap();
    let pids = [Pid::from(pid as usize)];
    sys.refresh_processes_specifics(
        ProcessesToUpdate::Some(&pids),
        true,
        ProcessRefreshKind::nothing().with_cpu().with_memory(),
    );

    let sys_pid = Pid::from(pid as usize);

    if let Some(process) = sys.process(sys_pid) {
        let ram = process.memory(); // bytes
        let cpu = process.cpu_usage() as f64; // percentage
        return Ok((ram, cpu));
    }

    Err(format!("Process not found: pid={}", pid))
}

/// Load a model via the llama-server Router API.
/// Uses the canonical `/models/load` endpoint with the `"model"` JSON field.
/// Model is identified by its display_name (matches the `--alias` registered in llama.cpp).
pub async fn load_model(host: &str, port: u16, model_id: &str) -> Result<(), String> {
    let host = clean_host(host);

    let props_url = format!("http://{}:{}/props", host, port);
    if let Ok(res) = HTTP_CLIENT.get(&props_url).send().await
        && res.status().is_success()
        && let Ok(json) = res.json::<serde_json::Value>().await
        && json.get("role").and_then(|r| r.as_str()) != Some("router")
    {
        return Err(
            "Server is not in router mode. Start with --models-max to enable router mode."
                .to_string(),
        );
    }

    let url = format!("http://{}:{}/models/load", host, port);
    let body = serde_json::json!({ "model": model_id });

    match HTTP_CLIENT.post(&url).json(&body).send().await {
        Ok(res) if res.status().is_success() => Ok(()),
        Ok(res) => {
            let status = res.status();
            let error = res
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            if error.contains("model is already running") {
                Ok(())
            } else {
                Err(format!("Server returned {}: {}", status, error))
            }
        }
        Err(e) => Err(format!("Failed to send request: {}", e)),
    }
}

/// List all models and their status from the llama-server Router API.
pub async fn list_models(
    host: &str,
    port: u16,
) -> Result<Vec<(String, String, Option<String>)>, String> {
    let host = clean_host(host);
    let url = format!("http://{}:{}/models", host, port);

    let res = HTTP_CLIENT
        .get(&url)
        .send()
        .await
        .map_err(|e| format!("Failed to list models: {}", e))?;

    if !res.status().is_success() {
        return Err(format!("Server returned error {}", res.status()));
    }

    let json: serde_json::Value = res
        .json()
        .await
        .map_err(|e| format!("Invalid JSON: {}", e))?;

    let mut results = Vec::new();
    if let Some(data) = json.get("data").and_then(|d| d.as_array()) {
        for model in data {
            let id = model
                .get("id")
                .and_then(|v| v.as_str())
                .unwrap_or_default()
                .to_string();
            // Status can be a string or an object with a "value" field
            let status = model
                .get("status")
                .and_then(|s| s.get("value").or(Some(s)))
                .and_then(|v| v.as_str())
                .unwrap_or("unloaded")
                .to_string();
            let path = model
                .get("path")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());

            results.push((id, status, path));
        }
    }

    Ok(results)
}

/// Unload a model via the llama-server Router API.
/// Uses the canonical `/models/unload` endpoint with the `"model"` JSON field.
pub async fn unload_model(host: &str, port: u16, model_id: &str) -> Result<(), String> {
    let host = clean_host(host);
    let url = format!("http://{}:{}/models/unload", host, port);
    let body = serde_json::json!({ "model": model_id });

    match HTTP_CLIENT.post(&url).json(&body).send().await {
        Ok(res) if res.status().is_success() => Ok(()),
        Ok(res) => {
            let status = res.status();
            let error = res
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            tracing::debug!(
                "Model unload failed (status {}, error: {}): this is expected if model was already unloaded",
                status,
                error
            );
            Ok(())
        }
        Err(e) => {
            tracing::debug!("Model unload request failed: {}", e);
            Ok(())
        }
    }
}