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

use std::sync::Arc;

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

use crate::AppState;
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"),
        "medium" => ModelId::new("whisper-medium-en"),
        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>,
}

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

#[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();

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

// --- 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-medium-en" => "Medium".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())
}