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

use std::sync::Arc;

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

use crate::commands::security::ensure_main_window;
use crate::AppState;
use lumotia_core::constants::WHISPER_SAMPLE_RATE;
use lumotia_core::hardware::{self, vulkan_loader_available, CpuFeatures};
use lumotia_core::model_registry::{self, Engine, LanguageSupport, ModelEntry};
use lumotia_core::types::{AudioSamples, ModelId, TranscriptionOptions};
#[cfg(feature = "whisper")]
use lumotia_transcription::load_whisper;
use lumotia_transcription::model_manager;
use lumotia_transcription::{load_parakeet, LocalEngine, Transcriber};

/// 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.to_string()),
    }
}

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

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<Box<dyn Transcriber + Send>, String> {
    let entry =
        model_registry::find_model(model_id).ok_or_else(|| format!("Unknown model: {model_id}"))?;

    match entry.engine {
        #[cfg(feature = "whisper")]
        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())
        }
        #[cfg(not(feature = "whisper"))]
        Engine::Whisper => Err(format!(
            "Whisper backend not compiled in this build (lumotia built without the \"whisper\" feature); \
             cannot load {model_id}"
        )),
        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,
    concurrent: Option<bool>,
) -> Result<(), String> {
    let model_id = ModelId::new(model_id.to_string());
    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(());
    }

    // Sequential-GPU guard (brief item A.1 #28): if the user opts out
    // of concurrent GPU residency, free the LLM before bringing the
    // transcription engine on. None / Some(true) leaves the LLM
    // untouched (legacy parallel behaviour, safe on multi-GB VRAM
    // setups). Inverse guard lives in commands::llm::load_llm_model.
    if concurrent == Some(false) && state.llm_engine.is_loaded() {
        state.llm_engine.unload().map_err(|e| e.to_string())?;
    }

    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);
                // Silent warm-up pass: feed one second of silence through
                // the freshly-loaded engine. Pre-allocates the Whisper
                // context window + warms GPU shader caches so the user's
                // first real transcription completes in ≤1.5× steady-state
                // latency instead of the ~4–5s cold-start documented in
                // ufal/whisper_streaming #96 and #135. Silence returns
                // empty segments — the *work* is the context allocation.
                let silence = AudioSamples::mono_16khz(vec![0.0_f32; WHISPER_SAMPLE_RATE as usize]);
                let options = TranscriptionOptions::default();
                match whisper_engine.transcribe_sync(&silence, &options) {
                    Ok(_) => tracing::info!(
                        target: "lumotia_startup",
                        "Whisper warm-up inference complete"
                    ),
                    Err(e) => tracing::warn!(
                        target: "lumotia_startup",
                        error = %e,
                        "Whisper warm-up inference failed"
                    ),
                }
            })
        })
        .await;

        match result {
            Ok(Ok(())) => tracing::info!(
                target: "lumotia_startup",
                "Whisper model pre-warmed successfully"
            ),
            Ok(Err(e)) => tracing::warn!(
                target: "lumotia_startup",
                error = %e,
                "Whisper pre-warm failed"
            ),
            Err(e) => tracing::error!(
                target: "lumotia_startup",
                error = %e,
                "Whisper pre-warm task panicked"
            ),
        }
    });
}

#[tauri::command]
pub async fn prewarm_default_model_cmd(
    window: tauri::WebviewWindow,
    state: tauri::State<'_, AppState>,
) -> Result<(), String> {
    ensure_main_window(&window)?;
    prewarm_default_model(state.whisper_engine.clone());
    Ok(())
}

#[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,
}

/// Compile-time target signalling used by `compose_accelerators`.
/// Split out so the pure-function behaviour is testable without `cfg!`
/// appearing in the test matrix.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum AcceleratorTarget {
    Macos,
    Other,
}

/// Pure helper: produce the `accelerators` list from the
/// whisper-compiled-in flag, runtime loader availability, and target
/// family. Always starts with `cpu`; appends the platform-appropriate
/// GPU name only when whisper is compiled in AND the Vulkan loader
/// resolves. RB-07 replaced a hard-coded `["cpu", "vulkan"]`.
fn compose_accelerators(
    whisper_enabled: bool,
    loader_available: bool,
    target: AcceleratorTarget,
) -> Vec<String> {
    let mut accelerators = vec!["cpu".into()];
    if whisper_enabled && loader_available {
        let gpu = match target {
            AcceleratorTarget::Macos => "metal",
            AcceleratorTarget::Other => "vulkan",
        };
        accelerators.push(gpu.into());
    }
    accelerators
}

/// Public wrapper around `compose_accelerators` that reads the real
/// `cfg(feature = "whisper")`, runtime loader probe, and target OS.
fn supported_accelerators() -> Vec<String> {
    let target = if cfg!(target_os = "macos") {
        AcceleratorTarget::Macos
    } else {
        AcceleratorTarget::Other
    };
    compose_accelerators(cfg!(feature = "whisper"), vulkan_loader_available(), target)
}

/// 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,
            };
        }
        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();

    // Accelerator list is now derived from the live build configuration
    // and runtime loader probe — see `compose_accelerators`. Parakeet
    // (ONNX) stays CPU-only; Whisper advertises `supports_gpu` only
    // when it was actually compiled in for this binary.
    let whisper_supports_gpu = cfg!(feature = "whisper");

    Ok(RuntimeCapabilities {
        accelerators: supported_accelerators(),
        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: whisper_supports_gpu,
                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(
    window: tauri::WebviewWindow,
    app: tauri::AppHandle,
    size: String,
) -> Result<String, String> {
    ensure_main_window(&window)?;
    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 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(
    window: tauri::WebviewWindow,
    state: tauri::State<'_, AppState>,
    size: String,
    concurrent: Option<bool>,
) -> Result<String, String> {
    ensure_main_window(&window)?;
    let id = whisper_model_id(&size);
    ensure_model_loaded(&state, "whisper", id.as_str(), concurrent).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(
    window: tauri::WebviewWindow,
    app: tauri::AppHandle,
    name: String,
) -> Result<String, String> {
    ensure_main_window(&window)?;
    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(
    window: tauri::WebviewWindow,
    state: tauri::State<'_, AppState>,
    name: String,
    concurrent: Option<bool>,
) -> Result<String, String> {
    ensure_main_window(&window)?;
    let id = parakeet_model_id(&name);
    ensure_model_loaded(&state, "parakeet", id.as_str(), concurrent).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())
}

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

    // RB-07: the accelerator list is now derived from compile flags +
    // runtime loader probe + target. These cover the permutations.
    //
    // Pre-fix behaviour was `vec!["cpu".into(), "vulkan".into()]`
    // regardless of whisper feature, loader availability, or platform —
    // so a macOS build falsely advertised "vulkan", and a no-whisper
    // build falsely advertised a GPU path with no engine to use it.

    #[test]
    fn cpu_only_when_whisper_disabled() {
        assert_eq!(
            compose_accelerators(false, true, AcceleratorTarget::Macos),
            vec!["cpu".to_string()]
        );
        assert_eq!(
            compose_accelerators(false, true, AcceleratorTarget::Other),
            vec!["cpu".to_string()]
        );
    }

    #[test]
    fn cpu_only_when_loader_missing() {
        assert_eq!(
            compose_accelerators(true, false, AcceleratorTarget::Macos),
            vec!["cpu".to_string()]
        );
        assert_eq!(
            compose_accelerators(true, false, AcceleratorTarget::Other),
            vec!["cpu".to_string()]
        );
    }

    #[test]
    fn macos_with_loader_advertises_metal() {
        assert_eq!(
            compose_accelerators(true, true, AcceleratorTarget::Macos),
            vec!["cpu".to_string(), "metal".to_string()]
        );
    }

    #[test]
    fn non_macos_with_loader_advertises_vulkan() {
        assert_eq!(
            compose_accelerators(true, true, AcceleratorTarget::Other),
            vec!["cpu".to_string(), "vulkan".to_string()]
        );
    }

    #[test]
    fn cpu_is_always_first_entry() {
        // Frontend relies on index-0 being the fallback; preserve that
        // contract regardless of which GPU extras are added.
        for target in [AcceleratorTarget::Macos, AcceleratorTarget::Other] {
            let full = compose_accelerators(true, true, target);
            assert_eq!(full.first(), Some(&"cpu".to_string()));
        }
    }
}