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

use tauri::{Emitter, State};

use crate::commands::power::PowerAssertion;
use crate::commands::security::ensure_main_window;
use crate::AppState;
use lumotia_ai_formatting::{llm_cleanup_text, LlmPromptPreset};
use lumotia_core::hardware;
use lumotia_llm::model_manager::{self, model_info};
use lumotia_llm::{ContentTags, LlmModelId};

#[derive(Debug, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LlmModelStatusDto {
    pub id: String,
    pub display_name: String,
    pub downloaded: bool,
    pub loaded: bool,
    pub size_bytes: u64,
    pub description: String,
    pub minimum_ram_bytes: u64,
    pub recommended_vram_bytes: Option<u64>,
}

fn parse_model_id(model_id: String) -> Result<LlmModelId, String> {
    model_id.parse()
}

#[tauri::command]
pub fn recommend_llm_tier() -> Result<String, String> {
    let profile = hardware::probe_system();
    let ram_bytes = profile.ram.0.saturating_mul(1024 * 1024);
    let vram_bytes = profile
        .gpu
        .map(|gpu| gpu.vram.0.saturating_mul(1024 * 1024));
    Ok(model_manager::recommend_tier(ram_bytes, vram_bytes)
        .as_str()
        .to_string())
}

#[tauri::command]
pub fn check_llm_model(
    state: State<'_, AppState>,
    model_id: String,
) -> Result<LlmModelStatusDto, String> {
    let id = parse_model_id(model_id)?;
    let info = model_info(id);
    let loaded_model_id = state.llm_engine.loaded_model_id();

    Ok(LlmModelStatusDto {
        id: info.id,
        display_name: info.display_name.to_string(),
        downloaded: model_manager::is_downloaded(id),
        loaded: loaded_model_id.as_deref() == Some(id.as_str()),
        size_bytes: info.size_bytes,
        description: info.description.to_string(),
        minimum_ram_bytes: info.minimum_ram_bytes,
        recommended_vram_bytes: info.recommended_vram_bytes,
    })
}

#[tauri::command]
pub async fn download_llm_model(
    window: tauri::WebviewWindow,
    app: tauri::AppHandle,
    model_id: String,
) -> Result<(), String> {
    ensure_main_window(&window)?;
    let id = parse_model_id(model_id)?;
    let app_clone = app.clone();
    model_manager::download_model(id, move |done, total| {
        let percent = if total > 0 {
            ((done as f64 / total as f64) * 100.0).round() as u8
        } else {
            0
        };
        let _ = app_clone.emit(
            "lumotia:llm-download-progress",
            serde_json::json!({
                "modelId": id.as_str(),
                "done": done,
                "total": total,
                "percent": percent,
            }),
        );
    })
    .await
    .map_err(|e| e.to_string())
}

#[tauri::command]
pub async fn load_llm_model(
    window: tauri::WebviewWindow,
    state: State<'_, AppState>,
    model_id: String,
    use_gpu: Option<bool>,
    concurrent: Option<bool>,
) -> Result<(), String> {
    ensure_main_window(&window)?;
    let id = parse_model_id(model_id)?;
    let path = model_manager::model_path(id);
    if !path.exists() {
        return Err("Model not downloaded — call download_llm_model first".to_string());
    }

    // Sequential-GPU guard (brief item A.1 #28): when the user has opted
    // out of concurrent GPU residency, free the transcription engines
    // before loading the LLM. Prevents VRAM OOM on tight-GPU setups.
    // concurrent=None or Some(true) preserves legacy parallel behaviour.
    if concurrent == Some(false) {
        state.whisper_engine.unload();
        state.parakeet_engine.unload();
    }

    let engine = state.llm_engine.clone();
    let use_gpu = use_gpu.unwrap_or(true);
    tokio::task::spawn_blocking(move || engine.load_model(id, &path, use_gpu))
        .await
        .map_err(|e| e.to_string())?
        .map_err(|e| e.to_string())
}

#[tauri::command]
pub fn unload_llm_model(
    window: tauri::WebviewWindow,
    state: State<'_, AppState>,
) -> Result<(), String> {
    ensure_main_window(&window)?;
    state.llm_engine.unload().map_err(|e| e.to_string())
}

#[tauri::command]
pub fn delete_llm_model(
    window: tauri::WebviewWindow,
    state: State<'_, AppState>,
    model_id: String,
) -> Result<(), String> {
    ensure_main_window(&window)?;
    let id = parse_model_id(model_id)?;
    if state.llm_engine.loaded_model_id().as_deref() == Some(id.as_str()) {
        state.llm_engine.unload().map_err(|e| e.to_string())?;
    }
    model_manager::delete_model(id).map_err(|e| e.to_string())
}

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

/// Diagnostic result for the Settings "Test LLM" button (brief item
/// B.1 #27). Classifies LLM setup failures into actionable categories
/// instead of surfacing a raw llama.cpp error string.
#[derive(Debug, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LlmTestResult {
    /// One of: "ready", "not-downloaded", "incomplete",
    /// "load-failed-vram", "load-failed-corrupt",
    /// "load-failed-permission", "load-failed-other".
    pub category: String,
    /// `true` when the LLM is healthy and usable after the test.
    pub ok: bool,
    /// One-line status copy for the Settings chip ("Qwen3.5 4B ready").
    pub message: String,
    /// Optional actionable next step ("Click Download", "Delete and
    /// re-download", "Pick a smaller tier"). Absent when the state is
    /// healthy.
    pub hint: Option<String>,
}

/// Best-effort LLM health check. Behaviour:
/// 1. Model not downloaded → reports `not-downloaded` with a
///    download hint.
/// 2. File present but size is ≤90% of expected → reports
///    `incomplete` (stalled download) with a re-download hint.
/// 3. Same model already loaded → returns `ready` without disturbing
///    the engine.
/// 4. Otherwise attempts `engine.load_model(...)` and classifies any
///    error string via `classify_llm_load_error` — VRAM exhaustion,
///    GGUF magic mismatch, filesystem permissions, or
///    everything-else. Success returns `ready`.
///
/// The point is that the user sees "Not enough GPU memory — pick a
/// smaller tier" rather than a raw C++ exception bubbled up from
/// llama.cpp. Mirrors OpenWhispr's "Test connection" UX for cloud
/// LLMs, adapted to Magnotia's local stack.
#[tauri::command]
pub async fn test_llm_model(
    window: tauri::WebviewWindow,
    state: State<'_, AppState>,
    model_id: String,
) -> Result<LlmTestResult, String> {
    ensure_main_window(&window)?;
    let id = parse_model_id(model_id)?;
    let info = model_info(id);
    let path = model_manager::model_path(id);

    if !path.exists() {
        return Ok(LlmTestResult {
            category: "not-downloaded".into(),
            ok: false,
            message: format!("{} is not downloaded.", info.display_name),
            hint: Some(format!(
                "Click Download in Settings → AI (~{} MB).",
                info.size_bytes / 1_000_000
            )),
        });
    }

    // Partial-download detection: llama-cpp-2 will segfault or panic
    // on a truncated GGUF rather than returning a clean error, so
    // catch it here before we attempt a load. 10% tolerance because
    // the expected size is rounded in model_manager.
    if let Ok(metadata) = std::fs::metadata(&path) {
        let actual = metadata.len();
        let minimum = info.size_bytes.saturating_sub(info.size_bytes / 10);
        if actual < minimum {
            return Ok(LlmTestResult {
                category: "incomplete".into(),
                ok: false,
                message: format!(
                    "{} file is incomplete ({} MB of expected {} MB).",
                    info.display_name,
                    actual / 1_000_000,
                    info.size_bytes / 1_000_000
                ),
                hint: Some("Delete and re-download from Settings → AI.".into()),
            });
        }
    }

    // Already loaded — no need to disturb the engine just to confirm.
    if state.llm_engine.loaded_model_id().as_deref() == Some(id.as_str()) {
        return Ok(LlmTestResult {
            category: "ready".into(),
            ok: true,
            message: format!("{} loaded and ready.", info.display_name),
            hint: None,
        });
    }

    // Not currently loaded. Attempt a real load (with GPU by default
    // — matches load_llm_model's default) and classify any failure.
    let engine = state.llm_engine.clone();
    let load_result = tokio::task::spawn_blocking(move || engine.load_model(id, &path, true))
        .await
        .map_err(|e| e.to_string())?;

    match load_result {
        Ok(()) => Ok(LlmTestResult {
            category: "ready".into(),
            ok: true,
            message: format!("{} loaded and ready.", info.display_name),
            hint: None,
        }),
        Err(err) => {
            let raw = err.to_string();
            let (category, hint) = classify_llm_load_error(&raw);
            Ok(LlmTestResult {
                category: category.into(),
                ok: false,
                message: format!("Load failed: {raw}"),
                hint: Some(hint.into()),
            })
        }
    }
}

/// Pure string classifier so the test_llm_model command stays
/// unit-testable without spinning up an actual LlmEngine. Order of
/// checks matters — permission errors can contain the word "failed"
/// too, so we check narrower categories before the catch-all.
fn classify_llm_load_error(raw: &str) -> (&'static str, &'static str) {
    let lower = raw.to_lowercase();
    if lower.contains("out of memory")
        || lower.contains("oom")
        || lower.contains("allocation failed")
        || lower.contains("vram")
        || lower.contains("cudamalloc")
    {
        (
            "load-failed-vram",
            "Not enough GPU memory. Pick a smaller tier in Settings → AI, or disable GPU acceleration (Advanced → GPU Tuning).",
        )
    } else if lower.contains("magic")
        || lower.contains("invalid gguf")
        || lower.contains("unsupported file format")
        || lower.contains("tensor shape")
    {
        (
            "load-failed-corrupt",
            "Model file appears corrupt or unsupported. Delete and re-download from Settings → AI.",
        )
    } else if lower.contains("permission denied") || lower.contains("access is denied") {
        (
            "load-failed-permission",
            "Permission denied reading the model file. Check ownership of ~/.magnotia/models/llm/.",
        )
    } else {
        (
            "load-failed-other",
            "Unexpected load error. See Settings → About → Diagnostics bundle.",
        )
    }
}

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

    #[test]
    fn classifies_vram_exhaustion() {
        let (category, hint) = classify_llm_load_error("cudaMalloc failed: out of memory");
        assert_eq!(category, "load-failed-vram");
        assert!(hint.contains("smaller tier"));
    }

    #[test]
    fn classifies_oom_alias() {
        let (category, _) = classify_llm_load_error("OOM while allocating 4096 MB");
        assert_eq!(category, "load-failed-vram");
    }

    #[test]
    fn classifies_generic_allocation_failure_as_vram() {
        let (category, _) = classify_llm_load_error("allocation failed at step 7");
        assert_eq!(category, "load-failed-vram");
    }

    #[test]
    fn classifies_gguf_magic_mismatch() {
        let (category, hint) = classify_llm_load_error("invalid gguf magic bytes");
        assert_eq!(category, "load-failed-corrupt");
        assert!(hint.contains("re-download"));
    }

    #[test]
    fn classifies_unsupported_format() {
        let (category, _) = classify_llm_load_error("Unsupported file format for model");
        assert_eq!(category, "load-failed-corrupt");
    }

    #[test]
    fn classifies_permission_denied() {
        let (category, hint) = classify_llm_load_error("os error 13: Permission denied");
        assert_eq!(category, "load-failed-permission");
        assert!(hint.contains("ownership"));
    }

    #[test]
    fn classifies_windows_access_denied() {
        let (category, _) = classify_llm_load_error("Access is denied. (os error 5)");
        assert_eq!(category, "load-failed-permission");
    }

    #[test]
    fn classifies_unknown_error_as_other() {
        let (category, _) = classify_llm_load_error("Quantum entanglement disrupted");
        assert_eq!(category, "load-failed-other");
    }
}

#[tauri::command]
pub async fn cleanup_transcript_text_cmd(
    window: tauri::WebviewWindow,
    state: State<'_, AppState>,
    transcript: String,
    profile_id: Option<String>,
    preset: Option<String>,
) -> Result<String, String> {
    ensure_main_window(&window)?;
    let resolved_profile_id =
        profile_id.unwrap_or_else(|| lumotia_storage::DEFAULT_PROFILE_ID.to_string());
    let profile_terms: Vec<String> =
        lumotia_storage::database::list_profile_terms(&state.db, &resolved_profile_id)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .map(|term| term.term)
            .collect();

    // Named preset (brief item B.1 #15): Email / Notes / Code shape the
    // output tone + structure without changing the translator-not-editor
    // contract. None or unknown → Default (no additional guidance).
    let resolved_preset = preset
        .as_deref()
        .map(LlmPromptPreset::parse)
        .unwrap_or(LlmPromptPreset::Default);

    let engine = state.llm_engine.clone();
    tokio::task::spawn_blocking(move || {
        // macOS: pin a power assertion for the duration of the LLM
        // generation so App Nap can't decide to throttle us mid-token.
        // No-op on every other OS. Item #9.
        let _power_guard = PowerAssertion::begin("magnotia LLM cleanup");
        llm_cleanup_text(&engine, &transcript, &profile_terms, resolved_preset)
    })
    .await
    .map_err(|e| e.to_string())?
    .map_err(|e| e.to_string())
}

/// Phase 9 LLM-powered content tags. On-demand from the History page;
/// never auto-runs. Heavy work (LlmEngine::extract_content_tags is
/// synchronous llama-cpp inference) is wrapped in spawn_blocking so it
/// does not stall the Tauri runtime, with the same App-Nap power
/// assertion the other LLM commands use.
#[tauri::command]
pub async fn extract_content_tags_cmd(
    state: State<'_, AppState>,
    transcript: String,
) -> Result<ContentTags, String> {
    if !state.llm_engine.is_loaded() {
        return Err("LLM not loaded. Download an AI model in Settings.".to_string());
    }
    let engine = state.llm_engine.clone();
    tokio::task::spawn_blocking(move || {
        let _power_guard = PowerAssertion::begin("magnotia LLM content-tag extraction");
        engine.extract_content_tags(&transcript)
    })
    .await
    .map_err(|e| e.to_string())?
    .map_err(|e| e.to_string())
}