Lumotia/src-tauri/src/commands/models.rs

use std::sync::Arc;

use serde::Serialize;
use tauri::Emitter;

use crate::AppState;
use kon_core::hardware::{self, CpuFeatures};
use kon_core::model_registry::{self, Engine, LanguageSupport, ModelEntry};
use kon_core::types::ModelId;
use kon_transcription::model_manager;
use kon_transcription::{load_parakeet, load_whisper, LocalEngine};

/// Map legacy size strings to ModelId.
fn whisper_model_id(size: &str) -> ModelId {
    match size.to_lowercase().as_str() {
        "tiny" => ModelId::new("whisper-tiny-en"),
        "base" => ModelId::new("whisper-base-en"),
        "small" => ModelId::new("whisper-small-en"),
        "distil-small" | "distilsmall" => ModelId::new("whisper-distil-small-en"),
        "medium" => ModelId::new("whisper-medium-en"),
        "distil-large" | "distil-large-v3" | "distillarge" => {
            ModelId::new("whisper-distil-large-v3")
        }
        other => ModelId::new(other),
    }
}

fn parakeet_model_id(name: &str) -> ModelId {
    match name {
        "ctc-int8" => ModelId::new("parakeet-ctc-0.6b-int8"),
        other => ModelId::new(other),
    }
}

fn engine_for_name(state: &AppState, engine_name: &str) -> Result<Arc<LocalEngine>, String> {
    match engine_name {
        "whisper" => Ok(state.whisper_engine.clone()),
        "parakeet" => Ok(state.parakeet_engine.clone()),
        other => Err(format!("Unknown engine: {other}")),
    }
}

fn language_support_info(language_support: LanguageSupport) -> LanguageSupportInfo {
    match language_support {
        LanguageSupport::EnglishOnly => LanguageSupportInfo {
            kind: "english-only".into(),
            language_count: 1,
        },
        LanguageSupport::Multilingual(count) => LanguageSupportInfo {
            kind: "multilingual".into(),
            language_count: count,
        },
    }
}

fn model_capability(
    entry: &'static ModelEntry,
    engine: &Arc<LocalEngine>,
) -> ModelRuntimeCapabilities {
    let loaded_model_id = engine.loaded_model_id();
    ModelRuntimeCapabilities {
        id: entry.id.to_string(),
        display_name: entry.display_name.to_string(),
        downloaded: model_manager::is_downloaded(&entry.id),
        loaded: loaded_model_id
            .as_ref()
            .map(|id| id == &entry.id)
            .unwrap_or(false),
        language_support: language_support_info(entry.languages),
    }
}

pub fn load_model_from_disk(
    model_id: &ModelId,
) -> Result<kon_transcription::SpeechBackend, String> {
    let entry =
        model_registry::find_model(model_id).ok_or_else(|| format!("Unknown model: {model_id}"))?;

    match entry.engine {
        Engine::Whisper => {
            let dir = model_manager::model_dir(model_id);
            let model_file = entry
                .files
                .first()
                .map(|file| dir.join(file.filename))
                .ok_or_else(|| format!("No files registered for model: {model_id}"))?;
            if !model_file.exists() {
                return Err(format!("Model not downloaded: {model_id}"));
            }
            load_whisper(&model_file).map_err(|e| e.to_string())
        }
        Engine::Parakeet => {
            let dir = model_manager::model_dir(model_id);
            if !dir.exists() {
                return Err(format!("Model not downloaded: {model_id}"));
            }
            load_parakeet(&dir).map_err(|e| e.to_string())
        }
        Engine::Moonshine => Err("Moonshine models are not yet supported in this build".into()),
    }
}

pub fn default_model_id_for_engine(engine: &str) -> ModelId {
    match engine {
        "parakeet" => ModelId::new("parakeet-ctc-0.6b-int8"),
        _ => ModelId::new("whisper-base-en"),
    }
}

pub async fn ensure_model_loaded(
    state: &AppState,
    engine_name: &str,
    model_id: &str,
) -> Result<(), String> {
    let model_id = ModelId::new(model_id);
    let entry = model_registry::find_model(&model_id)
        .ok_or_else(|| format!("Unknown model: {model_id}"))?;

    let expected_engine = match entry.engine {
        Engine::Whisper => "whisper",
        Engine::Parakeet => "parakeet",
        Engine::Moonshine => "moonshine",
    };
    if expected_engine != engine_name {
        return Err(format!(
            "Model {} belongs to {}, not {}",
            model_id, expected_engine, engine_name
        ));
    }

    if !model_manager::is_downloaded(&model_id) {
        return Err(format!("Model not downloaded: {model_id}"));
    }

    let engine = engine_for_name(state, engine_name)?;
    if engine
        .loaded_model_id()
        .as_ref()
        .map(|loaded| loaded == &model_id)
        .unwrap_or(false)
    {
        return Ok(());
    }

    let engine_clone = engine.clone();
    let model_id_clone = model_id.clone();
    tokio::task::spawn_blocking(move || {
        let model = load_model_from_disk(&model_id_clone)?;
        engine_clone.load(model, model_id_clone);
        Ok::<_, String>(())
    })
    .await
    .map_err(|e| e.to_string())??;

    Ok(())
}

/// Spawns a background task that loads the default Whisper model into memory
/// if it is already downloaded. Returns immediately — errors are logged, never panicked.
///
/// Call this once from app setup() after AppState is managed.
pub fn prewarm_default_model(whisper_engine: Arc<LocalEngine>) {
    let model_id = default_model_id_for_engine("whisper");

    if !model_manager::is_downloaded(&model_id) {
        return;
    }

    if whisper_engine
        .loaded_model_id()
        .as_ref()
        .map(|id| id == &model_id)
        .unwrap_or(false)
    {
        return;
    }

    tauri::async_runtime::spawn(async move {
        let result = tauri::async_runtime::spawn_blocking(move || {
            load_model_from_disk(&model_id).map(|model| {
                whisper_engine.load(model, model_id);
            })
        })
        .await;

        match result {
            Ok(Ok(())) => eprintln!("[startup] Whisper model pre-warmed successfully"),
            Ok(Err(e)) => eprintln!("[startup] Whisper pre-warm failed: {e}"),
            Err(e) => eprintln!("[startup] Whisper pre-warm task panicked: {e}"),
        }
    });
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RuntimeCapabilities {
    pub accelerators: Vec<String>,
    pub engines: Vec<EngineRuntimeCapabilities>,
    /// Which compute device Whisper / whisper.cpp actually booted against.
    /// Distinct from `accelerators` (which lists what the _build_ supports):
    /// a Vulkan-built binary on a CPU-only box reports accelerators=[cpu, vulkan]
    /// but activeComputeDevice.kind = "cpu" with a reason.
    pub active_compute_device: ActiveComputeDevice,
    /// Runtime-detected CPU feature flags. Surfaced so Settings can warn
    /// the user that performance will be poor without AVX2 / FMA.
    pub cpu_features: CpuFeaturesInfo,
    /// True when the detected hardware can sustain Whisper + LLM on the
    /// same GPU concurrently (≥16 GB VRAM). Item #28 gates a user-facing
    /// toggle on this.
    pub parallel_mode_available: bool,
}

/// Serialisable summary of whichever backend whisper.cpp / ggml wired
/// up on this boot. For MVP (Phase A.1) we derive this from
/// compile-time features + a runtime Vulkan loader probe; Phase A.2
/// wires `whisper_print_system_info` for the real answer.
#[derive(Serialize, Clone, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ActiveComputeDevice {
    /// One of "cuda" | "vulkan" | "metal" | "cpu".
    pub kind: String,
    /// Human-readable label, e.g. "GPU (Vulkan)" or "CPU (fallback)".
    pub label: String,
    /// Set only when we fell back from a richer backend, explains why
    /// (e.g. "Vulkan loader not found"). None on the happy path.
    pub reason: Option<String>,
}

#[derive(Serialize, Clone, Debug)]
#[serde(rename_all = "camelCase")]
pub struct CpuFeaturesInfo {
    pub avx2: bool,
    pub avx512f: bool,
    pub fma: bool,
    pub sse4_2: bool,
    pub neon: bool,
    /// Shortcut for Settings: false means we fall back to a slow path.
    pub has_ggml_baseline: bool,
}

impl From<CpuFeatures> for CpuFeaturesInfo {
    fn from(f: CpuFeatures) -> Self {
        Self {
            avx2: f.avx2,
            avx512f: f.avx512f,
            fma: f.fma,
            sse4_2: f.sse4_2,
            neon: f.neon,
            has_ggml_baseline: f.has_ggml_baseline(),
        }
    }
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct EngineRuntimeCapabilities {
    pub id: String,
    pub default_model_id: String,
    pub loaded_model_id: Option<String>,
    pub supports_gpu: bool,
    pub models: Vec<ModelRuntimeCapabilities>,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelRuntimeCapabilities {
    pub id: String,
    pub display_name: String,
    pub downloaded: bool,
    pub loaded: bool,
    pub language_support: LanguageSupportInfo,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LanguageSupportInfo {
    pub kind: String,
    pub language_count: u16,
}

/// Best-effort probe for the Vulkan loader shared library.
///
/// whisper.cpp's Vulkan backend refuses to initialise if `vulkan-1.dll`
/// (Windows) / `libvulkan.so.1` (Linux) / the MoltenVK bundle (macOS)
/// is missing at runtime. We probe via `libloading::Library::new`;
/// failure means whisper.cpp will silently drop back to CPU, and the
/// user deserves to know before they start wondering why a 14-second
/// clip takes two minutes.
fn vulkan_loader_available() -> bool {
    #[cfg(target_os = "linux")]
    let candidates: &[&str] = &["libvulkan.so.1", "libvulkan.so"];
    #[cfg(target_os = "windows")]
    let candidates: &[&str] = &["vulkan-1.dll"];
    #[cfg(target_os = "macos")]
    let candidates: &[&str] = &["libvulkan.1.dylib", "libMoltenVK.dylib"];
    #[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "macos")))]
    let candidates: &[&str] = &[];

    for name in candidates {
        // SAFETY: libloading::Library::new loads a shared library and returns
        // a handle that is dropped at the end of this iteration. No symbols
        // from it are called, so we just need the open-for-probe semantics.
        match unsafe { libloading::Library::new(*name) } {
            Ok(_lib) => return true,
            Err(_) => continue,
        }
    }
    false
}

/// Report which backend whisper.cpp was actually able to initialise
/// against. The whisper-rs build here is compiled with the `vulkan`
/// feature unconditionally; on macOS that's still Metal via MoltenVK,
/// on Linux/Windows it's Vulkan. If the Vulkan loader is missing
/// we surface the CPU fallback path explicitly so the UI can warn.
pub fn detect_active_compute_device() -> ActiveComputeDevice {
    #[cfg(target_os = "macos")]
    {
        if vulkan_loader_available() {
            return ActiveComputeDevice {
                kind: "metal".into(),
                label: "GPU (Metal via MoltenVK)".into(),
                reason: None,
            };
        }
        return ActiveComputeDevice {
            kind: "cpu".into(),
            label: "CPU (fallback)".into(),
            reason: Some(
                "MoltenVK / Vulkan loader not found — install the Vulkan SDK runtime."
                    .into(),
            ),
        };
    }

    #[cfg(not(target_os = "macos"))]
    {
        if vulkan_loader_available() {
            return ActiveComputeDevice {
                kind: "vulkan".into(),
                label: "GPU (Vulkan)".into(),
                reason: None,
            };
        }
        return ActiveComputeDevice {
            kind: "cpu".into(),
            label: "CPU (fallback)".into(),
            reason: Some(
                "Vulkan loader not found — install the Vulkan runtime (Windows) or \
                 libvulkan1 (Linux) to enable GPU acceleration."
                    .into(),
            ),
        };
    }
}

#[derive(Serialize, Clone, Debug)]
#[serde(rename_all = "camelCase", tag = "kind", content = "message")]
pub enum RuntimeWarning {
    /// CPU lacks AVX2 / FMA on x86_64 — ggml fallback path will be
    /// dramatically slower. User should install a non-AVX2 build or
    /// accept the hit.
    Avx2Missing(String),
    /// Vulkan loader is missing at runtime. Emitted once at startup
    /// when `active_compute_device.reason` includes a loader-missing
    /// message.
    VulkanLoaderMissing(String),
    /// CUDA driver mismatch forced a fallback (future-proofed for
    /// when we add CUDA-detect; not emitted today).
    #[allow(dead_code)]
    CudaFallback(String),
}

/// Emit any runtime warnings the frontend should surface as a banner.
/// Called once at setup() after `prewarm_default_model`.
pub fn emit_runtime_warnings(app: &tauri::AppHandle) {
    let cpu_features = hardware::probe_cpu_features();
    let device = detect_active_compute_device();

    if !cpu_features.has_ggml_baseline() {
        let _ = app.emit(
            "runtime-warning",
            RuntimeWarning::Avx2Missing(
                "Your CPU is missing AVX2/FMA. Whisper and the local LLM will run on a \
                 dramatically slower fallback path — expect 5–10× the latency of a \
                 2015-or-newer CPU."
                    .into(),
            ),
        );
    }

    if device.kind == "cpu" {
        if let Some(reason) = device.reason.as_ref() {
            let _ = app.emit(
                "runtime-warning",
                RuntimeWarning::VulkanLoaderMissing(reason.clone()),
            );
        }
    }
}

#[tauri::command]
pub fn get_runtime_capabilities(
    state: tauri::State<'_, AppState>,
) -> Result<RuntimeCapabilities, String> {
    let whisper = state.whisper_engine.clone();
    let parakeet = state.parakeet_engine.clone();

    let whisper_models = model_registry::all_models()
        .iter()
        .filter(|entry| entry.engine == Engine::Whisper)
        .map(|entry| model_capability(entry, &whisper))
        .collect();
    let parakeet_models = model_registry::all_models()
        .iter()
        .filter(|entry| entry.engine == Engine::Parakeet)
        .map(|entry| model_capability(entry, &parakeet))
        .collect();

    let active_compute_device = detect_active_compute_device();
    let cpu_features: CpuFeaturesInfo = hardware::probe_cpu_features().into();

    // Kon's desktop build links transcribe-rs with the `whisper-vulkan`
    // feature unconditionally (see crates/transcription/Cargo.toml), so
    // whisper.cpp boots with Vulkan backend on any machine with a Vulkan
    // loader + ICD. Parakeet (ONNX) still runs on CPU. Reflect both.
    Ok(RuntimeCapabilities {
        accelerators: vec!["cpu".into(), "vulkan".into()],
        engines: vec![
            EngineRuntimeCapabilities {
                id: "whisper".into(),
                default_model_id: default_model_id_for_engine("whisper").to_string(),
                loaded_model_id: whisper
                    .loaded_model_id()
                    .map(|model_id| model_id.to_string()),
                supports_gpu: true,
                models: whisper_models,
            },
            EngineRuntimeCapabilities {
                id: "parakeet".into(),
                default_model_id: default_model_id_for_engine("parakeet").to_string(),
                loaded_model_id: parakeet
                    .loaded_model_id()
                    .map(|model_id| model_id.to_string()),
                supports_gpu: false,
                models: parakeet_models,
            },
        ],
        active_compute_device,
        cpu_features,
        // Phase A.1 ships detection of the VRAM budget only via
        // hardware::probe_gpu; that currently returns None, so we
        // default to sequential. When Phase A.4 (GpuGuard) and the
        // real probe_gpu land, flip this based on detected VRAM ≥ 16 GB.
        parallel_mode_available: false,
    })
}

// --- Whisper model commands ---

#[tauri::command]
pub async fn download_model(app: tauri::AppHandle, size: String) -> Result<String, String> {
    let id = whisper_model_id(&size);
    let app_clone = app.clone();
    model_manager::download(&id, move |progress| {
        let _ = app_clone.emit("model-download-progress", &progress);
    })
    .await
    .map_err(|e| e.to_string())?;
    Ok(format!("Model {} downloaded", size))
}

#[tauri::command]
pub fn check_model(size: String) -> Result<bool, String> {
    let id = whisper_model_id(&size);
    Ok(model_manager::is_downloaded(&id))
}

#[tauri::command]
pub fn list_models() -> Result<Vec<String>, String> {
    let ids = model_manager::list_downloaded();
    Ok(ids
        .into_iter()
        .filter(|id| id.as_str().starts_with("whisper-"))
        .map(|id| match id.as_str() {
            "whisper-tiny-en" => "Tiny".to_string(),
            "whisper-base-en" => "Base".to_string(),
            "whisper-small-en" => "Small".to_string(),
            "whisper-distil-small-en" => "Distil-S".to_string(),
            "whisper-medium-en" => "Medium".to_string(),
            "whisper-distil-large-v3" => "Distil-L".to_string(),
            other => other.to_string(),
        })
        .collect())
}

#[tauri::command]
pub async fn load_model(state: tauri::State<'_, AppState>, size: String) -> Result<String, String> {
    let id = whisper_model_id(&size);
    ensure_model_loaded(&state, "whisper", id.as_str()).await?;
    Ok(format!("Model {} loaded", size))
}

#[tauri::command]
pub fn check_engine(state: tauri::State<AppState>) -> Result<bool, String> {
    Ok(state.whisper_engine.is_loaded())
}

// --- Parakeet model commands ---

#[tauri::command]
pub async fn download_parakeet_model(
    app: tauri::AppHandle,
    name: String,
) -> Result<String, String> {
    let id = parakeet_model_id(&name);
    let app_clone = app.clone();
    model_manager::download(&id, move |progress| {
        let _ = app_clone.emit("parakeet-download-progress", &progress);
    })
    .await
    .map_err(|e| e.to_string())?;
    Ok(format!("Parakeet model {} downloaded", name))
}

#[tauri::command]
pub fn check_parakeet_model(name: String) -> Result<bool, String> {
    let id = parakeet_model_id(&name);
    Ok(model_manager::is_downloaded(&id))
}

#[tauri::command]
pub fn list_parakeet_models() -> Result<Vec<String>, String> {
    let ids = model_manager::list_downloaded();
    Ok(ids
        .into_iter()
        .filter(|id| id.as_str().starts_with("parakeet-"))
        .map(|id| match id.as_str() {
            "parakeet-ctc-0.6b-int8" => "ctc-int8".to_string(),
            other => other.to_string(),
        })
        .collect())
}

#[tauri::command]
pub async fn load_parakeet_model(
    state: tauri::State<'_, AppState>,
    name: String,
) -> Result<String, String> {
    let id = parakeet_model_id(&name);
    ensure_model_loaded(&state, "parakeet", id.as_str()).await?;
    Ok(format!("Parakeet model {} loaded", name))
}

#[tauri::command]
pub fn check_parakeet_engine(state: tauri::State<'_, AppState>) -> Result<bool, String> {
    Ok(state.parakeet_engine.is_loaded())
}