Lumotia/crates/llm/src/lib.rs

use std::num::NonZeroU32;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use encoding_rs::UTF_8;
use llama_cpp_2::context::params::LlamaContextParams;
use llama_cpp_2::llama_backend::LlamaBackend;
use llama_cpp_2::llama_batch::LlamaBatch;
use llama_cpp_2::model::params::LlamaModelParams;
use llama_cpp_2::model::{AddBos, LlamaChatMessage, LlamaChatTemplate, LlamaModel};
use llama_cpp_2::sampling::LlamaSampler;
use lumotia_core::tuning::{inference_thread_count, Workload};
use serde::{Deserialize, Serialize};

pub mod grammars;
pub mod model_manager;
pub mod prompts;

pub use grammars::CONTENT_TAGS_GRAMMAR;
pub use model_manager::{recommend_tier, LlmModelId, LlmModelInfo};
pub use prompts::{is_valid_intent, ContentTags, CONTENT_TAGS_SYSTEM, INTENT_CLOSED_SET};

const DEFAULT_CONTEXT_TOKENS: u32 = 4096;
const MAX_CONTEXT_TOKENS: u32 = 8192;
const CONTEXT_RESERVE_TOKENS: u32 = 64;
const GENERATION_SEED: u32 = 0;

/// Maximum number of tasks returned by the rule-based fallback extractor.
/// Caps output to avoid wall-of-text dumps when the transcript is dense.
const MAX_RULE_BASED_TASKS: usize = 10;

/// Indicates which extraction path produced the task list.
/// Propagated to callers so the UI can label rule-based results accordingly.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TaskExtractionSource {
    /// Tasks extracted by the local LLM.
    Llm,
    /// LLM path failed; tasks extracted by the rule-based regex fallback.
    RuleBased,
}

#[derive(Debug, thiserror::Error)]
pub enum EngineError {
    #[error("LLM not loaded. Download an AI model in Settings.")]
    NotLoaded,
    #[error("LLM load failed: {0}")]
    LoadFailed(String),
    #[error(
        "Another LLM load is already in flight; refusing to start a parallel load \
         or modify engine state mid-load."
    )]
    AlreadyLoading,
    #[error(
        "prompt too long: {prompt_tokens} prompt tokens exceed the {available_prompt_tokens}-token prompt budget for an {context_window}-token context with {max_tokens} reserved response tokens"
    )]
    PromptTooLong {
        prompt_tokens: usize,
        max_tokens: u32,
        available_prompt_tokens: u32,
        context_window: u32,
    },
    #[error("inference failed: {0}")]
    Inference(String),
    #[error("model output not valid JSON: {0}")]
    InvalidJson(String),
}

#[derive(Debug, Clone)]
pub struct GenerationConfig {
    pub max_tokens: u32,
    pub temperature: f32,
    pub stop_sequences: Vec<String>,
    pub grammar: Option<String>,
}

impl Default for GenerationConfig {
    fn default() -> Self {
        Self {
            max_tokens: 1024,
            temperature: 0.0,
            stop_sequences: Vec::new(),
            grammar: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct LoadedModelState {
    pub model_id: String,
    pub model_path: String,
    pub use_gpu: bool,
}

#[derive(Default)]
struct LlmState {
    backend: Option<Arc<LlamaBackend>>,
    model: Option<Arc<LlamaModel>>,
    loaded: Option<LoadedModelState>,
}

#[derive(Clone, Default)]
pub struct LlmEngine {
    inner: Arc<Mutex<LlmState>>,
    /// Flag held for the duration of a model load. The std::sync::Mutex
    /// covers cheap state mutations (~microseconds); the multi-second
    /// `LlamaModel::load_from_file` call now runs *outside* the mutex,
    /// so polls like `is_loaded()` / `loaded_model_id()` (called from
    /// sync Tauri handlers without `spawn_blocking`) don't park tokio
    /// worker threads on a slow C++ FFI call. This Atomic also doubles
    /// as a TOCTOU guard: two concurrent `load_model` invocations on
    /// the same engine will not both reach the heavy load — the second
    /// returns `EngineError::AlreadyLoading`.
    loading: Arc<AtomicBool>,
}

/// RAII guard that clears the `loading` flag on drop, including on
/// panic / early-return. Prevents the engine getting stuck in a
/// "permanently loading" state if a load fails midway.
struct LoadingGuard {
    flag: Arc<AtomicBool>,
}

impl Drop for LoadingGuard {
    fn drop(&mut self) {
        self.flag.store(false, Ordering::Release);
    }
}

impl LlmEngine {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn load(&self, model_path: &Path) -> Result<(), EngineError> {
        self.load_model(LlmModelId::default_tier(), model_path, true)
    }

    // instrument: the load is multi-second (`LlamaBackend::init` +
    // mmap + GPU layer init). Tagging events with `model_id` and
    // `use_gpu` lets the operator separate the GPU sequential-guard
    // logs and llama-backend init lines from the LLM transcription
    // pipeline by structured field rather than by adjacency.
    #[tracing::instrument(skip_all, fields(model_id = %model_id.as_str(), use_gpu = use_gpu))]
    pub fn load_model(
        &self,
        model_id: LlmModelId,
        model_path: &Path,
        use_gpu: bool,
    ) -> Result<(), EngineError> {
        self.load_model_with(model_id, model_path, use_gpu, |backend, path, params| {
            LlamaModel::load_from_file(backend, path, params)
                .map_err(|e| EngineError::LoadFailed(format!("model load: {e}")))
        })
    }

    /// Core load implementation with a swappable file-loader closure.
    /// Production callers use `load_model`, which delegates here with
    /// the real `LlamaModel::load_from_file`. Tests inject a sleepy /
    /// counting closure to exercise the locking discipline without
    /// pulling a real GGUF off disk.
    ///
    /// Locking discipline (the whole point of this function):
    ///   1. Take the mutex briefly to compare against the currently
    ///      loaded triple — if it matches, return early. No-op fast path.
    ///   2. CAS the `loading` flag from false → true. If another load is
    ///      already in flight, refuse with `AlreadyLoading` rather than
    ///      starting a parallel one. A `LoadingGuard` ensures the flag
    ///      is cleared on every exit path including panic.
    ///   3. Take the mutex briefly to drop the OLD model Arc (frees its
    ///      VRAM via `llama_free_model`) before the new load begins.
    ///      The backend Arc is preserved — `LlamaBackend::init()` is a
    ///      one-shot per process (an `AtomicBool` in llama-cpp-2 enforces
    ///      `BackendAlreadyInitialized` on a second call), so we must
    ///      never drop the backend while the process keeps running.
    ///      Note: `is_loaded()` reports false during the swap window —
    ///      that is the correct semantics. Callers wanting "model X is
    ///      loaded" must check `loaded_model_id()` against their target.
    ///   4. Initialise the backend if absent (first-ever load only) and
    ///      run the slow `load_from_file` call — both OUTSIDE the mutex.
    ///   5. Take the mutex briefly to install the new backend (if just
    ///      initialised) and the new model Arc.
    fn load_model_with<F>(
        &self,
        model_id: LlmModelId,
        model_path: &Path,
        use_gpu: bool,
        loader: F,
    ) -> Result<(), EngineError>
    where
        F: FnOnce(&LlamaBackend, &Path, &LlamaModelParams) -> Result<LlamaModel, EngineError>,
    {
        // Step 1: short crit section — already-loaded fast path.
        {
            let guard = self.inner.lock().unwrap();
            if let Some(loaded) = &guard.loaded {
                if loaded.model_id == model_id.as_str()
                    && loaded.model_path == model_path.display().to_string()
                    && loaded.use_gpu == use_gpu
                {
                    return Ok(());
                }
            }
        }

        // Step 2: claim the loading slot. Refuse if a parallel load is
        // already mid-flight rather than starting a second slow load
        // and silently overwriting the first.
        if self
            .loading
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_err()
        {
            return Err(EngineError::AlreadyLoading);
        }
        let _loading_guard = LoadingGuard {
            flag: Arc::clone(&self.loading),
        };

        // Step 3: short crit section — drop the OLD model so its VRAM is
        // released BEFORE we allocate the new one. Without this, a swap
        // briefly holds two models resident (Lifecycle-1: an
        // ~17 GB Q4 27B swap on a 24 GB card OOMs even though either
        // model fits alone). Keep the backend Arc — see locking notes.
        let existing_backend = {
            let mut guard = self.inner.lock().unwrap();
            guard.model = None;
            guard.loaded = None;
            guard.backend.clone()
        };

        // Step 4: heavy work OUTSIDE the mutex. `is_loaded()` and
        // `loaded_model_id()` can be polled freely here without parking
        // tokio worker threads.
        let backend = match existing_backend {
            Some(existing) => existing,
            None => Arc::new(
                LlamaBackend::init()
                    .map_err(|e| EngineError::LoadFailed(format!("backend init: {e}")))?,
            ),
        };

        let gpu_layers = if use_gpu { u32::MAX } else { 0 };
        let params = LlamaModelParams::default().with_n_gpu_layers(gpu_layers);
        let model = loader(&backend, model_path, &params)?;

        // Step 5: short crit section — install the new state.
        {
            let mut guard = self.inner.lock().unwrap();
            guard.backend = Some(backend);
            guard.model = Some(Arc::new(model));
            guard.loaded = Some(LoadedModelState {
                model_id: model_id.as_str().to_string(),
                model_path: model_path.display().to_string(),
                use_gpu,
            });
        }
        // `_loading_guard` drops here and clears the flag.
        Ok(())
    }

    pub fn unload(&self) -> Result<(), EngineError> {
        // Refuse to unload mid-load. Without this check, `load_model_with`
        // is mid-flight (it has cleared `model` / `loaded` in step 3 and
        // is about to install new state in step 5); a concurrent unload
        // would do nothing (the state is already None), return Ok, and
        // then the load's step 5 silently overwrites — the caller saw
        // unload success but the engine ends up loaded. Phase B.7 audit
        // residual (2026-05-14): the load-vs-load TOCTOU was closed by
        // `AlreadyLoading` in cde985d but the unload-vs-load race was
        // left open. Same flag covers both directions.
        if self.is_loading() {
            return Err(EngineError::AlreadyLoading);
        }
        let mut guard = self.inner.lock().unwrap();
        guard.model = None;
        // Backend is process-singleton (llama-cpp-2 enforces this via
        // `LLAMA_BACKEND_INITIALIZED`). Dropping the Arc here would call
        // `llama_backend_free` and a subsequent `init` would succeed, but
        // we keep it resident to avoid the init/free churn on every
        // load/unload cycle.
        guard.loaded = None;
        Ok(())
    }

    /// True iff a model load is currently in flight. Exposed for tests
    /// and frontends that want to render a "loading…" state without
    /// polling `is_loaded()` (which returns false during a swap).
    pub fn is_loading(&self) -> bool {
        self.loading.load(Ordering::Acquire)
    }

    /// Test-only harness: runs `op` while holding the same locking
    /// discipline as `load_model_with` (loading flag claimed, model
    /// state cleared, slow op runs OUTSIDE the inner mutex, new state
    /// installed at the end). Used by the regression test to verify
    /// that `is_loaded()` / `loaded_model_id()` don't block on the
    /// slow section. Not part of the public API.
    #[cfg(test)]
    pub(crate) fn __test_run_with_lock_discipline<F>(&self, op: F) -> Result<(), EngineError>
    where
        F: FnOnce(),
    {
        if self
            .loading
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_err()
        {
            return Err(EngineError::AlreadyLoading);
        }
        let _loading_guard = LoadingGuard {
            flag: Arc::clone(&self.loading),
        };
        {
            let mut guard = self.inner.lock().unwrap();
            guard.model = None;
            guard.loaded = None;
        }
        op();
        Ok(())
    }

    pub fn is_loaded(&self) -> bool {
        self.inner.lock().unwrap().model.is_some()
    }

    pub fn loaded_model(&self) -> Option<LoadedModelState> {
        self.inner.lock().unwrap().loaded.clone()
    }

    pub fn loaded_model_id(&self) -> Option<String> {
        self.loaded_model().map(|loaded| loaded.model_id)
    }

    pub fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String, EngineError> {
        let (backend, model) = self.loaded_handles()?;
        let prompt_tokens = model
            .str_to_token(prompt, AddBos::Never)
            .map_err(|e| EngineError::Inference(format!("tokenize: {e}")))?;
        if prompt_tokens.is_empty() {
            return Ok(String::new());
        }

        let n_ctx = preflight_context_window(prompt_tokens.len(), config.max_tokens)?;
        let use_gpu = self.loaded_model().map(|s| s.use_gpu).unwrap_or(false);
        let gpu_layers = if use_gpu { u32::MAX } else { 0u32 };
        // Trivially true today (gpu_layers = u32::MAX when use_gpu), but the
        // explicit comparison documents intent. True residency observability
        // (parsing llama.cpp's "offloaded N/M layers" log) is tracked as a
        // follow-up in docs/superpowers/specs/2026-05-09-battery-gpu-aware-
        // thread-tuning-design.md (§ Out of scope).
        let gpu_offloaded = use_gpu && gpu_layers >= model.n_layer();
        let thread_count =
            i32::try_from(inference_thread_count(Workload::Llm, gpu_offloaded)).unwrap_or(4);
        let ctx_params = LlamaContextParams::default()
            .with_n_ctx(Some(
                NonZeroU32::new(n_ctx).expect("n_ctx must be non-zero"),
            ))
            .with_n_batch(prompt_tokens.len().max(512).min(n_ctx as usize) as u32)
            .with_n_ubatch(prompt_tokens.len().max(512).min(n_ctx as usize) as u32)
            .with_n_threads(thread_count)
            .with_n_threads_batch(thread_count);
        let mut ctx = model
            .new_context(&backend, ctx_params)
            .map_err(|e| EngineError::Inference(format!("context: {e}")))?;

        let mut batch = LlamaBatch::new(prompt_tokens.len().max(1), 1);
        for (index, token) in prompt_tokens.iter().enumerate() {
            batch
                .add(*token, index as i32, &[0], index + 1 == prompt_tokens.len())
                .map_err(|e| EngineError::Inference(format!("batch add: {e}")))?;
        }
        ctx.decode(&mut batch)
            .map_err(|e| EngineError::Inference(format!("prefill decode: {e}")))?;

        let mut sampler = self.build_sampler(&model, config)?;
        let mut decoder = UTF_8.new_decoder();
        let mut generated = String::new();
        let mut cursor = prompt_tokens.len() as i32;

        for _ in 0..config.max_tokens {
            let next = sampler.sample(&ctx, batch.n_tokens() - 1);
            if model.is_eog_token(next) || next == model.token_eos() {
                break;
            }

            let piece = model
                .token_to_piece(next, &mut decoder, true, None)
                .map_err(|e| EngineError::Inference(format!("detokenize: {e}")))?;
            generated.push_str(&piece);
            sampler.accept(next);

            if config.grammar.is_none() && json_envelope_complete(&generated) {
                break;
            }

            if let Some(stop_index) = first_stop_index(&generated, &config.stop_sequences) {
                generated.truncate(stop_index);
                break;
            }

            batch.clear();
            batch
                .add(next, cursor, &[0], true)
                .map_err(|e| EngineError::Inference(format!("sample batch: {e}")))?;
            cursor += 1;
            ctx.decode(&mut batch)
                .map_err(|e| EngineError::Inference(format!("sample decode: {e}")))?;
        }

        Ok(generated.trim().to_string())
    }

    pub fn cleanup_text(
        &self,
        system_prompt: &str,
        transcript: &str,
    ) -> Result<String, EngineError> {
        if transcript.trim().is_empty() {
            return Ok(String::new());
        }
        let model = self.loaded_model_arc()?;
        let prompt =
            render_chat_prompt(&model, &[("system", system_prompt), ("user", transcript)])?;
        self.generate(
            &prompt,
            &GenerationConfig {
                max_tokens: 1024,
                temperature: 0.0,
                stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
                grammar: None,
            },
        )
    }

    pub fn decompose_task(&self, task_text: &str) -> Result<Vec<String>, EngineError> {
        self.decompose_task_with_feedback(task_text, &[])
    }

    /// Same as `decompose_task` but allows callers to pass recent HITL
    /// feedback rows so the system prompt gets conditioned on the
    /// user's preferred decomposition style. The `examples` vec is
    /// rendered into a few-shot block appended to the base system
    /// prompt by `prompts::build_conditioned_system_prompt`.
    ///
    /// Callers should pass most-recent-first; older examples still
    /// participate but weigh less because of their position in the
    /// prompt. Empty slice keeps behaviour identical to `decompose_task`.
    pub fn decompose_task_with_feedback(
        &self,
        task_text: &str,
        examples: &[prompts::FeedbackExample],
    ) -> Result<Vec<String>, EngineError> {
        let model = self.loaded_model_arc()?;
        let system =
            prompts::build_conditioned_system_prompt(prompts::DECOMPOSE_TASK_SYSTEM, examples);
        let prompt = render_chat_prompt(
            &model,
            &[
                ("system", system.as_str()),
                ("user", &format!("Task: {task_text}")),
            ],
        )?;
        let raw = self.generate(
            &prompt,
            &GenerationConfig {
                max_tokens: 512,
                temperature: 0.0,
                stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
                grammar: Some(grammars::TASK_ARRAY_GRAMMAR.to_string()),
            },
        )?;
        parse_string_array(&raw)
    }

    pub fn extract_tasks(&self, transcript: &str) -> Result<Vec<String>, EngineError> {
        self.extract_tasks_with_feedback(transcript, &[])
    }

    /// Phase 9 content-tag extraction. Emits a single (topic, intent)
    /// pair as JSON. Truncates to the trailing 2000 chars of the
    /// transcript so the prompt budget stays well under any model's
    /// context window. Determinism is enforced by temperature 0.0;
    /// the parsed intent is re-validated with `is_valid_intent` and
    /// invalid JSON bubbles as `InvalidJson` so the frontend toasts a
    /// clear error.
    pub fn extract_content_tags(
        &self,
        transcript: &str,
    ) -> Result<prompts::ContentTags, EngineError> {
        if transcript.trim().is_empty() {
            return Err(EngineError::Inference("empty transcript".into()));
        }

        // Truncate to the last 2000 chars on a UTF-8 char boundary so
        // we don't slice through a multi-byte sequence.
        const MAX_CHARS: usize = 2000;
        let tail = if transcript.len() > MAX_CHARS {
            let mut adj = transcript.len() - MAX_CHARS;
            while adj < transcript.len() && !transcript.is_char_boundary(adj) {
                adj += 1;
            }
            &transcript[adj..]
        } else {
            transcript
        };

        let model = self.loaded_model_arc()?;
        let prompt = render_chat_prompt(
            &model,
            &[
                ("system", prompts::CONTENT_TAGS_SYSTEM),
                ("user", &format!("Transcript:\n{tail}")),
            ],
        )?;
        let raw = self.generate(
            &prompt,
            &GenerationConfig {
                max_tokens: 96,
                temperature: 0.0,
                stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
                grammar: None,
            },
        )?;

        let tags: prompts::ContentTags = parse_json_payload(&raw)?;
        if !prompts::is_valid_intent(&tags.intent) {
            return Err(EngineError::InvalidJson(format!(
                "intent out of closed set: {}",
                tags.intent,
            )));
        }
        Ok(tags)
    }

    /// Feedback-conditioned variant of `extract_tasks`. See
    /// `decompose_task_with_feedback` for the `examples` semantics.
    pub fn extract_tasks_with_feedback(
        &self,
        transcript: &str,
        examples: &[prompts::FeedbackExample],
    ) -> Result<Vec<String>, EngineError> {
        if transcript.trim().is_empty() {
            return Ok(Vec::new());
        }

        let model = self.loaded_model_arc()?;
        let system =
            prompts::build_conditioned_system_prompt(prompts::EXTRACT_TASKS_SYSTEM, examples);
        let prompt = render_chat_prompt(
            &model,
            &[
                ("system", system.as_str()),
                ("user", &format!("Transcript:\n{transcript}")),
            ],
        )?;
        let raw = self.generate(
            &prompt,
            &GenerationConfig {
                max_tokens: 768,
                temperature: 0.0,
                stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
                grammar: Some(grammars::OPTIONAL_TASK_ARRAY_GRAMMAR.to_string()),
            },
        )?;
        parse_string_array(&raw)
    }

    /// Wrapper around [`extract_tasks_with_feedback`] that NEVER returns
    /// an error: if the LLM path fails for any reason the rule-based
    /// extractor fires as a safety net, satisfying the data-loss contract
    /// documented in `docs/release/v0.1-known-limitations.md`.
    ///
    /// Returns `(tasks, source)` where `source` tells the caller which
    /// path produced the results so the UI can label them.
    pub fn extract_tasks_with_fallback(
        &self,
        transcript: &str,
        examples: &[prompts::FeedbackExample],
    ) -> (Vec<String>, TaskExtractionSource) {
        match self.extract_tasks_with_feedback(transcript, examples) {
            Ok(tasks) => (tasks, TaskExtractionSource::Llm),
            Err(err) => {
                tracing::warn!(
                    "LLM task extraction failed; using rule-based fallback: {}",
                    err
                );
                let tasks = rule_based_extract_tasks(transcript);
                (tasks, TaskExtractionSource::RuleBased)
            }
        }
    }

    fn loaded_handles(&self) -> Result<(Arc<LlamaBackend>, Arc<LlamaModel>), EngineError> {
        let guard = self.inner.lock().unwrap();
        let backend = guard.backend.clone().ok_or(EngineError::NotLoaded)?;
        let model = guard.model.clone().ok_or(EngineError::NotLoaded)?;
        Ok((backend, model))
    }

    fn loaded_model_arc(&self) -> Result<Arc<LlamaModel>, EngineError> {
        self.loaded_handles().map(|(_, model)| model)
    }

    fn build_sampler(
        &self,
        model: &LlamaModel,
        config: &GenerationConfig,
    ) -> Result<LlamaSampler, EngineError> {
        let mut samplers = Vec::new();

        if let Some(grammar) = &config.grammar {
            samplers.push(
                LlamaSampler::grammar(model, grammar, "root")
                    .map_err(|e| EngineError::Inference(format!("grammar: {e}")))?,
            );
        }

        if config.temperature <= f32::EPSILON {
            samplers.push(LlamaSampler::greedy());
        } else {
            samplers.push(LlamaSampler::temp(config.temperature));
            samplers.push(LlamaSampler::dist(GENERATION_SEED));
        }

        Ok(if samplers.len() == 1 {
            samplers.remove(0)
        } else {
            LlamaSampler::chain_simple(samplers)
        })
    }
}

fn context_window_size(prompt_tokens: usize, max_tokens: u32) -> u32 {
    let required = prompt_tokens
        .saturating_add(max_tokens as usize)
        .saturating_add(CONTEXT_RESERVE_TOKENS as usize);
    DEFAULT_CONTEXT_TOKENS.max(required.min(MAX_CONTEXT_TOKENS as usize) as u32)
}

fn preflight_context_window(prompt_tokens: usize, max_tokens: u32) -> Result<u32, EngineError> {
    let required = prompt_tokens
        .saturating_add(max_tokens as usize)
        .saturating_add(CONTEXT_RESERVE_TOKENS as usize);
    if required > MAX_CONTEXT_TOKENS as usize {
        let available_prompt_tokens =
            MAX_CONTEXT_TOKENS.saturating_sub(max_tokens.saturating_add(CONTEXT_RESERVE_TOKENS));
        return Err(EngineError::PromptTooLong {
            prompt_tokens,
            max_tokens,
            available_prompt_tokens,
            context_window: MAX_CONTEXT_TOKENS,
        });
    }

    Ok(context_window_size(prompt_tokens, max_tokens))
}

fn first_stop_index(text: &str, stop_sequences: &[String]) -> Option<usize> {
    stop_sequences
        .iter()
        .filter(|stop| !stop.is_empty())
        .filter_map(|stop| text.find(stop))
        .min()
}

fn json_envelope_complete(text: &str) -> bool {
    extract_json_envelope(text) == Some(text.trim())
}

fn extract_json_envelope(text: &str) -> Option<&str> {
    // Phase B.9 audit residual (2026-05-14): strip the leading
    // `<think>…</think>` reasoning block before scanning. Qwen-style
    // models emit non-empty reasoning when thinking mode is on, and
    // the reasoning can contain JSON-looking literals (e.g.
    // "the answer should be {\"x\":1}") or unbalanced braces ("I wonder
    // about {..."). The naive "find the first '{' or '['" extractor
    // would then either return the wrong envelope or pollute the
    // brace-stack and return None. We split on the FIRST `</think>` —
    // anything before it is reasoning, anything after is the answer
    // proper. Falls back to the whole text when no `</think>` is
    // present (covers non-reasoning models and the empty-thinking
    // case already covered by `extract_json_envelope_skips_qwen_thinking_prefix`).
    let scan_region = text
        .split_once("</think>")
        .map(|(_, rest)| rest)
        .unwrap_or(text);

    let start = scan_region
        .char_indices()
        .find_map(|(idx, ch)| (ch == '{' || ch == '[').then_some(idx))?;
    let mut chars = scan_region[start..].char_indices();
    let (_, first) = chars.next()?;

    let mut stack = vec![match first {
        '{' => '}',
        '[' => ']',
        _ => unreachable!(),
    }];
    let mut in_string = false;
    let mut escaped = false;

    for (offset, ch) in chars {
        if in_string {
            if escaped {
                escaped = false;
            } else if ch == '\\' {
                escaped = true;
            } else if ch == '"' {
                in_string = false;
            }
            continue;
        }

        match ch {
            '"' => in_string = true,
            '{' => stack.push('}'),
            '[' => stack.push(']'),
            '}' | ']' => {
                if stack.pop() != Some(ch) {
                    return None;
                }
                if stack.is_empty() {
                    let end = start + offset + ch.len_utf8();
                    return Some(&scan_region[start..end]);
                }
            }
            _ => {}
        }
    }

    None
}

fn parse_json_payload<T: for<'de> Deserialize<'de>>(raw: &str) -> Result<T, EngineError> {
    let payload = extract_json_envelope(raw).unwrap_or(raw.trim());
    serde_json::from_str(payload).map_err(|e| EngineError::InvalidJson(format!("{e}: raw={raw:?}")))
}

fn render_chat_prompt(
    model: &LlamaModel,
    messages: &[(&str, &str)],
) -> Result<String, EngineError> {
    let chat_messages = messages
        .iter()
        .map(|(role, content)| {
            LlamaChatMessage::new((*role).to_string(), (*content).to_string())
                .map_err(|e| EngineError::Inference(format!("chat message: {e}")))
        })
        .collect::<Result<Vec<_>, _>>()?;

    match model.chat_template(None) {
        Ok(template) => model
            .apply_chat_template(&template, &chat_messages, true)
            .map_err(|e| EngineError::Inference(format!("chat template apply: {e}"))),
        Err(err) => {
            tracing::warn!("model chat template unavailable, falling back to ChatML: {err}");
            let template = LlamaChatTemplate::new("chatml")
                .map_err(|e| EngineError::Inference(format!("chatml template: {e}")))?;
            model
                .apply_chat_template(&template, &chat_messages, true)
                .map_err(|e| EngineError::Inference(format!("chatml template apply: {e}")))
        }
    }
}

fn parse_string_array(raw: &str) -> Result<Vec<String>, EngineError> {
    let parsed = serde_json::from_str::<Vec<String>>(raw.trim())
        .map_err(|e| EngineError::InvalidJson(format!("{e} in: {raw:?}")))?;

    let mut seen = std::collections::HashSet::new();
    let normalized = parsed
        .into_iter()
        .map(|item| item.trim().to_string())
        .filter(|item| !item.is_empty())
        .filter(|item| seen.insert(item.to_lowercase()))
        .collect();

    Ok(normalized)
}

/// Rule-based task extractor used as the safety net when the LLM extraction
/// path fails. Per `docs/release/v0.1-known-limitations.md`, task extraction
/// must NEVER return zero tasks just because the LLM failed.
///
/// Heuristic: split on sentence boundaries (`. ? ! \n`), keep sentences that
/// begin with (or contain near the start) an imperative-style cue. Trim,
/// dedupe, cap at [`MAX_RULE_BASED_TASKS`] to avoid wall-of-text dumps.
pub fn rule_based_extract_tasks(transcript: &str) -> Vec<String> {
    // Filler words that may precede the real imperative start.
    const FILLER: &[&str] = &["and ", "so ", "then ", "also ", "well ", "okay ", "ok "];

    // Phrase-level cues (checked against the lowercased sentence start).
    const PHRASE_CUES: &[&str] = &[
        "i need to ",
        "i should ",
        "i have to ",
        "i must ",
        "need to ",
        "got to ",
        "have to ",
        "must ",
        "let me ",
        "let's ",
        "lets ",
        "remember to ",
        "don't forget to ",
        "dont forget to ",
        "don't forget ",
        "dont forget ",
        "make sure to ",
        "make sure i ",
        "todo:",
        "to-do:",
        "task:",
    ];

    // Bare imperative verbs expected at the start of a sentence.
    const IMPERATIVE_VERBS: &[&str] = &[
        "send",
        "write",
        "call",
        "email",
        "fix",
        "update",
        "review",
        "check",
        "finish",
        "schedule",
        "book",
        "order",
        "buy",
        "ask",
        "follow up",
        "followup",
        "create",
        "add",
        "remove",
        "delete",
        "submit",
        "upload",
        "download",
        "install",
        "configure",
        "test",
        "deploy",
        "merge",
        "close",
        "open",
        "share",
        "contact",
        "reach out",
        "prepare",
        "draft",
        "complete",
        "reply",
        "respond",
    ];

    // Split on sentence-terminating punctuation and newlines.
    let sentences: Vec<&str> = transcript.split(['.', '?', '!', '\n']).collect();

    let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();
    let mut results: Vec<String> = Vec::new();

    for raw in sentences {
        let trimmed = raw.trim();
        if trimmed.is_empty() {
            continue;
        }

        // Build a lowercase version for matching, stripping leading filler.
        let mut lc = trimmed.to_lowercase();
        for filler in FILLER {
            if lc.starts_with(filler) {
                lc = lc[filler.len()..].trim_start().to_string();
                break;
            }
        }

        let is_task = PHRASE_CUES.iter().any(|cue| lc.starts_with(cue))
            || IMPERATIVE_VERBS.iter().any(|verb| {
                lc.starts_with(verb)
                    && lc
                        .as_bytes()
                        .get(verb.len())
                        .map(|&b| b == b' ' || b == b',')
                        .unwrap_or(true)
            });

        if is_task {
            let key = lc.clone();
            if seen.insert(key) {
                results.push(trimmed.to_string());
                if results.len() >= MAX_RULE_BASED_TASKS {
                    break;
                }
            }
        }
    }

    results
}

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

    #[test]
    fn generate_fails_when_not_loaded() {
        let engine = LlmEngine::new();
        let err = engine
            .generate("hello", &GenerationConfig::default())
            .unwrap_err();
        assert!(matches!(err, EngineError::NotLoaded));
    }

    #[test]
    fn decompose_returns_error_when_not_loaded() {
        let engine = LlmEngine::new();
        assert!(!engine.is_loaded());
        let result = engine.decompose_task("Write a blog post");
        assert!(matches!(result, Err(EngineError::NotLoaded)));
    }

    #[test]
    fn default_creates_unloaded_engine() {
        let engine = LlmEngine::default();
        assert!(!engine.is_loaded());
    }

    #[test]
    fn engine_is_clone_and_shares_state() {
        let engine = LlmEngine::new();
        let clone = engine.clone();
        assert!(!clone.is_loaded());
    }

    #[test]
    fn parse_string_array_trims_and_dedupes() {
        let parsed = parse_string_array(r#"[" Buy milk ", "buy milk", "Call plumber"]"#).unwrap();
        assert_eq!(parsed, vec!["Buy milk", "Call plumber"]);
    }

    #[test]
    fn first_stop_index_finds_earliest_match() {
        let text = "hello<|im_end|>trailing";
        let index = first_stop_index(text, &["<|im_end|>".into(), "zzz".into()]);
        assert_eq!(index, Some(5));
    }

    #[test]
    fn json_envelope_complete_detects_finished_object() {
        assert!(json_envelope_complete(
            r#"{"topic":"meeting","intent":"planning"}"#
        ));
    }

    #[test]
    fn json_envelope_complete_detects_finished_array() {
        assert!(json_envelope_complete(r#"["Call plumber","Buy milk"]"#));
    }

    #[test]
    fn json_envelope_complete_ignores_braces_inside_strings() {
        assert!(!json_envelope_complete(r#"{"topic":"literal } brace""#));
    }

    #[test]
    fn json_envelope_complete_rejects_prefixes_and_trailing_text() {
        assert!(!json_envelope_complete(r#"{"topic":"meeting""#));
        assert!(!json_envelope_complete(r#"{"topic":"meeting"} extra"#));
    }

    #[test]
    fn extract_json_envelope_skips_qwen_thinking_prefix() {
        let raw =
            "<think>\n\n</think>\n\n{\"topic\":\"grant-application\",\"intent\":\"planning\"}";
        assert_eq!(
            extract_json_envelope(raw),
            Some("{\"topic\":\"grant-application\",\"intent\":\"planning\"}"),
        );
    }

    #[test]
    fn extract_json_envelope_handles_arrays_and_trailing_stop_text() {
        assert_eq!(
            extract_json_envelope("prefix [\"Call plumber\",\"Buy milk\"]<|im_end|>"),
            Some("[\"Call plumber\",\"Buy milk\"]"),
        );
    }

    /// Phase B.9 audit regression (2026-05-14). The original
    /// `extract_json_envelope_skips_qwen_thinking_prefix` test only
    /// covered an EMPTY `<think></think>` block. Qwen-style reasoning
    /// is typically non-empty and can contain JSON-looking literals
    /// (the model thinking out loud about what shape it should emit).
    /// The naive "first '{' wins" extractor mis-identified the
    /// reasoning's literal as the answer envelope and returned it,
    /// skipping the actual answer that followed `</think>`.
    ///
    /// Post-fix the extractor strips the leading `<think>…</think>`
    /// block before scanning, so the reasoning's literal cannot
    /// poison the result.
    #[test]
    fn extract_json_envelope_skips_thinking_block_with_json_looking_content() {
        let raw = "<think>The answer should look like {\"topic\":\"reasoning-example\",\"intent\":\"capture\"} \
                   based on the schema.</think>{\"topic\":\"real-answer\",\"intent\":\"planning\"}";
        assert_eq!(
            extract_json_envelope(raw),
            Some("{\"topic\":\"real-answer\",\"intent\":\"planning\"}"),
        );
    }

    /// Phase B.9 audit regression (2026-05-14). If the reasoning block
    /// contains UNBALANCED braces (e.g. the model writes "I wonder
    /// about {..." inside `<think>…</think>`), the pre-strip extractor
    /// would start its stack on that unbalanced `{`, never find a
    /// matching `}`, and continue past `</think>` polluting the stack
    /// with the real answer's braces — ultimately returning None and
    /// losing the answer entirely. Stripping the reasoning block first
    /// makes both cases moot.
    #[test]
    fn extract_json_envelope_survives_unbalanced_braces_in_thinking() {
        let raw = "<think>I wonder about {something unfinished here</think>{\"topic\":\"recovery\",\"intent\":\"capture\"}";
        assert_eq!(
            extract_json_envelope(raw),
            Some("{\"topic\":\"recovery\",\"intent\":\"capture\"}"),
        );
    }

    #[test]
    fn prompt_preflight_rejects_oversized_prompt_tokens() {
        let err = preflight_context_window(7_105, 1_024).unwrap_err();
        assert!(matches!(
            err,
            EngineError::PromptTooLong {
                prompt_tokens: 7_105,
                max_tokens: 1_024,
                available_prompt_tokens: 7_104,
                context_window: MAX_CONTEXT_TOKENS,
            }
        ));
    }

    #[test]
    fn prompt_preflight_keeps_prompts_within_budget() {
        let n_ctx = preflight_context_window(7_104, 1_024).unwrap();
        assert_eq!(n_ctx, MAX_CONTEXT_TOKENS);
    }

    /// Race-3 regression. The inner `std::sync::Mutex` MUST NOT be held
    /// across the slow `LlamaModel::load_from_file` call: sync Tauri
    /// command handlers like `get_llm_status`, `check_llm_model`,
    /// `delete_llm_model`, and `test_llm_model` call `is_loaded()` /
    /// `loaded_model_id()` from tokio worker threads without
    /// `spawn_blocking`. If the lock is held for the duration of a
    /// 5-15 s load, parallel status polls from the frontend park the
    /// tokio executor and the whole UI deadlocks.
    ///
    /// Pre-fix this test FAILS — both probes time out because the load
    /// holds the mutex. Post-fix it PASSES — probes return in ≤50 ms.
    #[test]
    fn is_loaded_does_not_block_on_slow_load() {
        use std::sync::mpsc;
        use std::thread;
        use std::time::{Duration, Instant};

        let engine = LlmEngine::new();
        let load_started = Arc::new(std::sync::Barrier::new(2));
        let release_load = Arc::new(std::sync::Barrier::new(2));

        let engine_for_loader = engine.clone();
        let load_started_for_loader = Arc::clone(&load_started);
        let release_load_for_loader = Arc::clone(&release_load);

        let loader_handle = thread::spawn(move || {
            engine_for_loader
                .__test_run_with_lock_discipline(|| {
                    // Signal the probe thread that the load is mid-flight
                    // (loading flag claimed, inner mutex released).
                    load_started_for_loader.wait();
                    // Wait until the probe thread says it's done so the
                    // load's "duration" is bounded by the probes.
                    release_load_for_loader.wait();
                })
                .unwrap();
        });

        // Wait until the loader is inside its slow section.
        load_started.wait();

        // Now probe `is_loaded()` and `loaded_model_id()` from this
        // thread. They MUST return without contending on the lock.
        let probe_deadline = Duration::from_millis(50);
        let (tx, rx) = mpsc::channel();
        let engine_for_probe = engine.clone();
        let probe_handle = thread::spawn(move || {
            let start = Instant::now();
            let loaded = engine_for_probe.is_loaded();
            let id = engine_for_probe.loaded_model_id();
            let loading = engine_for_probe.is_loading();
            let elapsed = start.elapsed();
            tx.send((loaded, id, loading, elapsed)).unwrap();
        });

        let result = rx
            .recv_timeout(probe_deadline)
            .expect("is_loaded / loaded_model_id probe must return within 50 ms");
        let (loaded, id, loading, elapsed) = result;

        assert!(
            !loaded,
            "is_loaded() should report false while a load is in flight"
        );
        assert_eq!(id, None, "loaded_model_id() should be None mid-load");
        assert!(loading, "is_loading() should report true mid-load");
        assert!(
            elapsed < probe_deadline,
            "probe took {elapsed:?}, expected < {probe_deadline:?}"
        );

        probe_handle.join().unwrap();
        release_load.wait();
        loader_handle.join().unwrap();

        // After the load completes the flag clears.
        assert!(!engine.is_loading());
    }

    /// Race-3 / Race-4 — concurrent load attempts must not both reach
    /// the heavy work. The second caller should be told `AlreadyLoading`
    /// rather than starting a parallel load that silently overwrites
    /// the first.
    #[test]
    fn second_concurrent_load_is_refused() {
        use std::thread;

        let engine = LlmEngine::new();
        let load_started = Arc::new(std::sync::Barrier::new(2));
        let release_load = Arc::new(std::sync::Barrier::new(2));

        let engine_for_loader = engine.clone();
        let load_started_for_loader = Arc::clone(&load_started);
        let release_load_for_loader = Arc::clone(&release_load);

        let loader_handle = thread::spawn(move || {
            engine_for_loader
                .__test_run_with_lock_discipline(|| {
                    load_started_for_loader.wait();
                    release_load_for_loader.wait();
                })
                .unwrap();
        });

        load_started.wait();

        // Second concurrent attempt MUST be refused, not parallel-load.
        let second = engine.__test_run_with_lock_discipline(|| {
            panic!("second concurrent load should never reach its op");
        });
        assert!(matches!(second, Err(EngineError::AlreadyLoading)));

        release_load.wait();
        loader_handle.join().unwrap();

        // After the first load completes, a fresh attempt is allowed.
        assert!(engine.__test_run_with_lock_discipline(|| {}).is_ok());
    }

    /// Phase B.7 audit regression (2026-05-14). The cde985d fix
    /// introduced the `loading` AtomicBool to refuse a second concurrent
    /// load, but left `unload()` blind to the flag. A concurrent unload
    /// during a load's slow window observed `model == None` (the load's
    /// step 3 had already cleared state), no-op-cleared the same nulls,
    /// returned Ok — and then the load's step 5 silently installed the
    /// new state. The caller saw unload-success but the engine ended up
    /// loaded.
    ///
    /// Post-fix: an unload mid-load is refused with
    /// `EngineError::AlreadyLoading`. The caller can retry once the
    /// load completes (signalled by `is_loading() == false`).
    #[test]
    fn unload_during_load_is_refused() {
        use std::thread;

        let engine = LlmEngine::new();
        let load_started = Arc::new(std::sync::Barrier::new(2));
        let release_load = Arc::new(std::sync::Barrier::new(2));

        let engine_for_loader = engine.clone();
        let load_started_for_loader = Arc::clone(&load_started);
        let release_load_for_loader = Arc::clone(&release_load);

        let loader_handle = thread::spawn(move || {
            engine_for_loader
                .__test_run_with_lock_discipline(|| {
                    load_started_for_loader.wait();
                    release_load_for_loader.wait();
                })
                .unwrap();
        });

        // Wait until the loader is mid-slow-section (loading flag claimed,
        // engine state cleared).
        load_started.wait();

        // Unload while the load is in flight MUST be refused, not silently
        // no-op'd and then overwritten by the load's install step.
        let result = engine.unload();
        assert!(
            matches!(result, Err(EngineError::AlreadyLoading)),
            "unload during a load must surface AlreadyLoading, got {result:?}"
        );

        release_load.wait();
        loader_handle.join().unwrap();

        // After the load completes the flag clears and unload succeeds.
        assert!(!engine.is_loading());
        engine
            .unload()
            .expect("unload after load completes must succeed");
    }

    // ── rule_based_extract_tasks ──────────────────────────────────────────

    #[test]
    fn rule_based_extract_finds_explicit_imperatives() {
        let t = "I need to send Sarah the report tomorrow. Don't forget the slide deck.";
        let tasks = rule_based_extract_tasks(t);
        assert_eq!(tasks.len(), 2, "expected 2 tasks, got: {tasks:?}");
        assert!(
            tasks[0].to_lowercase().contains("send sarah"),
            "first task should mention 'send sarah': {tasks:?}"
        );
        assert!(
            tasks[1].to_lowercase().contains("slide deck"),
            "second task should mention 'slide deck': {tasks:?}"
        );
    }

    #[test]
    fn rule_based_extract_caps_at_max() {
        let t = "Send email. Write doc. Call client. Fix bug. Update spec. Review PR. Check tests. Finish report. Schedule meeting. Book hotel. Order parts. Buy supplies.";
        let tasks = rule_based_extract_tasks(t);
        assert!(
            tasks.len() <= MAX_RULE_BASED_TASKS,
            "expected at most {MAX_RULE_BASED_TASKS} tasks, got {}",
            tasks.len()
        );
    }

    #[test]
    fn rule_based_extract_returns_empty_for_no_imperatives() {
        let t = "The weather is lovely today. The garden looks nice.";
        let tasks = rule_based_extract_tasks(t);
        assert_eq!(tasks.len(), 0, "expected 0 tasks, got: {tasks:?}");
    }

    #[test]
    fn rule_based_extract_dedupes_repeated_sentences() {
        let t = "I need to send the report. I need to send the report.";
        let tasks = rule_based_extract_tasks(t);
        assert_eq!(
            tasks.len(),
            1,
            "expected 1 deduplicated task, got: {tasks:?}"
        );
    }
}