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;

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

    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)
}

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