mahbot 0.4.0

//! Semantic (vector/embedding) search for archived tickets.
//!
//! # Why this exists
//!
//! This module provides a local Candle + GGUF-based embedding model that converts
//! ticket descriptions into dense vectors. These vectors enable **semantic search** —
//! finding archived tickets by *conceptual similarity* (e.g. `"authentication bug"`
//! matching `"login issue"`) — which pure FTS keyword search alone cannot do.
//!
//! The embedding model uses **jinaai/jina-embeddings-v5-text-nano-retrieval**
//! (Q4_K_M GGUF, ~150 MB), a EuroBERT architecture (LLaMA-style encoder without
//! causal masking) loaded via the Candle framework. The model and tokenizer are
//! downloaded on first use and cached in `~/.mahbot/models/`.
//!
//! The embedder is loaded **lazily on first `embed_query()` / `embed_document()` call**
//! and cached in a global [`RwLock`]; embedding is computed at **ticket creation time**
//! (to pre-compute a vector for future archived search via `embed_document()`) and at
//! **search-query time** (to vectorize the query for `SearchArchivedTicketsTool` via
//! `embed_query()`). Both paths gracefully degrade: if the model files haven't been
//! downloaded yet (or download fails), both functions return `None` and the caller
//! falls back to FTS-only search. A background retry loop downloads the model with
//! exponential backoff, making the embedder available without requiring a restart.
//!
//! ## Product decision
//!
//! **Do not propose removing this module or its dependencies without explicit
//! user approval.** It is a deliberate product decision, not accidental bloat
//! or dead code.

use crate::util::UnwrapPoison;
use crate::util::model_state::{AtomicModelState, ModelLoadGuard, ModelState};
use anyhow::{Context, Result, anyhow};
use candle_core::quantized::{QMatMul, gguf_file};
use candle_core::{DType, Device, Tensor};
use candle_nn::rotary_emb::rope_slow;
use candle_nn::{Embedding, Module};
use std::collections::HashMap;
use std::fs::File;
use std::path::Path;
use std::sync::RwLock;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokenizers::Tokenizer;
use tracing::{debug, info, warn};

// ── Constants ────────────────────────────────────────────────────────

/// Maximum sequence length for the embedding model (jina-embeddings-v5 supports 8192).
const MAX_SEQ_LEN: usize = 8192;

/// RoPE base frequency for the model.
const ROPE_FREQ_BASE: f32 = 1_000_000.0;

/// Timeout for model file download (10 minutes for ~150 MB).
const MODEL_DOWNLOAD_TIMEOUT: Duration = Duration::from_mins(10);

/// Default pad token ID for this model (tokenizer's eos_token_id = 128001).
const DEFAULT_PAD_ID: u32 = 128_001;

// ── Model URLs ───────────────────────────────────────────────────────

/// HuggingFace URL for the quantized GGUF model file.
const MODEL_URL: &str = "https://huggingface.co/jinaai/jina-embeddings-v5-text-nano-retrieval/resolve/main/v5-nano-retrieval-Q4_K_M.gguf";

/// SHA256 checksum of the model file (verified at download time).
const MODEL_SHA256: &str = "f50822244ba0c7a348c5455b99bb8a0afd182511e8a816888c5dc65d972e51d5";

/// HuggingFace URL for the tokenizer file.
const TOKENIZER_URL: &str = "https://huggingface.co/jinaai/jina-embeddings-v5-text-nano-retrieval/resolve/main/tokenizer.json";

/// SHA256 checksum of the tokenizer file (verified at download time).
const TOKENIZER_SHA256: &str = "98d4a1d32152d6cedf85b5e88f3b205106dca1fe72aaab34e0ac13c238421069";

// ── Model filenames ──────────────────────────────────────────────────

/// Filename of the quantized GGUF model file.
///
/// Extracted as a constant to avoid drift between the synchronous cache-load, background
/// download, and test paths when upgrading the model.
const MODEL_FILENAME: &str = "v5-nano-retrieval-Q4_K_M.gguf";

/// Filename of the tokenizer JSON file.
///
/// Must be kept in sync with [`MODEL_FILENAME`] — both files come from the same model repo.
const TOKENIZER_FILENAME: &str = "embed_tokenizer.json";

// ── Global state ─────────────────────────────────────────────────────

/// Global embedder singleton, wrapped in an Option for graceful degradation.
static GLOBAL_EMBEDDER: RwLock<Option<Embedder>> = RwLock::new(None);

/// Atomic state tracker to coordinate lazy initialization.
///
/// Embedder state machine (shared wrapper extracted in
/// [`crate::util::model_state::ModelState`]).
///
/// # States
///
/// | Value | Name     | Meaning                                          |
/// |-------|----------|--------------------------------------------------|
/// | 0     | UNINIT   | Embedder not loaded yet; first [`embed_query()`] or     |
/// |       |          | [`embed_document()`] call triggers initialization via   |
/// |       |          | [`ensure_embedder()`].                                  |
/// | 1     | LOADING  | Initialization in progress (sync cache load or   |
/// |       |          | background download with retries).               |
/// | 2     | READY    | A usable [`Embedder`] instance is available.     |
/// | 3     | FAILED   | Embedder initialization failed terminally;               |
/// |       |          | [`ensure_embedder()`] returns `false` immediately.       |
/// |       |          | A restart is required to retry.                         |
///
/// # Transitions
///
/// | From       | To         | Trigger                                              |
/// |------------|------------|------------------------------------------------------|
/// | `UNINIT`   | `LOADING`  | Atomic CAS in [`ensure_embedder()`] — exactly one    |
/// |            |            | caller becomes the initializer.                      |
/// | `LOADING`  | `READY`    | [`set_embedder_ready()`] after a successful load     |
/// |            |            | (sync cache hit, cached re-load, or fresh download). |
/// | `LOADING`  | `FAILED`   | **Code-level error:** freshly downloaded,          |
/// |            |            | SHA256-verified files fail to load (bug            |
/// |            |            | in [`Embedder::load()`]).                           |
/// |            |            | [`download_retry_loop()`] transitions state to      |
/// |            |            | [`ModelState::Failed`] and returns.                       |
/// | `LOADING`  | `FAILED`   | **HTTP client build failure:** reqwest client      |
/// |            |            | construction fails in [`download_retry_loop()`]     |
/// |            |            | (e.g. TLS backend unavailable).                     |
/// |            |            | [`download_retry_loop()`] transitions state to      |
/// |            |            | [`ModelState::Failed`] and returns.                       |
/// | `LOADING`  | `FAILED`   | **Background task panic:** [`ModelLoadGuard`]'s    |
/// |            |            | [`Drop`] guard transitions state to [`ModelState::Failed`]|
/// |            |            | when [`download_retry_loop()`] is dropped without   |
/// |            |            | reaching a terminal state (e.g. a panic in Candle   |
/// |            |            | or tokenizers).                                      |
/// | `LOADING`  | `FAILED`   | **No tokio runtime:** [`ensure_embedder()`] checks   |
/// |            |            | [`tokio::runtime::Handle::try_current()`] and        |
/// |            |            | transitions to [`ModelState::Failed`] if no tokio runtime  |
/// |            |            | is available (cannot spawn background download).      |
///
/// Once STATE reaches `READY` or `FAILED` it stays there for the lifetime of the process.
static STATE: AtomicModelState = AtomicModelState::new(ModelState::Uninit);

/// Atomically store a ready [`Embedder`] and transition state to [`ModelState::Ready`].
///
/// This bundles the two operations (RwLock write + state transition) into a single
/// call so no future code path can do one without the other.
#[inline]
fn set_embedder_ready(emb: Embedder) {
    *GLOBAL_EMBEDDER.write().unwrap_poison() = Some(emb);
    STATE.store(ModelState::Ready, Ordering::Release);
}

/// Try to initialize the embedder (sync load from cache or spawn background download).
///
/// Called on every [`embed_query()`] / [`embed_document()`] invocation. Returns `true` if the embedder is
/// ready, `false` if it's still loading or permanently unavailable.
///
/// # Graceful degradation
///
/// The sync cache-load path runs synchronously. If no tokio runtime is available to
/// spawn the background download task, [`STATE`] transitions from [`ModelState::Loading`] to
/// [`ModelState::Failed`] so subsequent calls return `false` immediately. Panics during sync
/// load (e.g., Candle-internal bugs or OOM) propagate to the caller — OOM aborts the
/// process regardless, and Candle panics are code bugs that should fail loud and fast.
fn ensure_embedder() -> bool {
    match STATE.load(Ordering::Acquire) {
        ModelState::Ready => return true,
        ModelState::Uninit => {} // proceed to initialize
        _ => return false,       // ModelState::Loading, ModelState::Failed, or unexpected state
    }

    if !STATE.transition(ModelState::Uninit, ModelState::Loading) {
        return false;
    }

    // ── Sync load from cache ──────────────────────────────────────────
    // No panic-catching needed: OOM aborts the process regardless, and
    // Candle-internal panics are code bugs that should fail loud and fast.
    let models_dir = crate::util::models_dir()
        .expect("CONFIG must be initialized before embedder initialization");
    let model_path = models_dir.join(MODEL_FILENAME);
    let tokenizer_path = models_dir.join(TOKENIZER_FILENAME);

    std::fs::create_dir_all(&models_dir).ok();

    if model_path.exists()
        && tokenizer_path.exists()
        && let Some(emb) = try_load_embedder(&model_path, &tokenizer_path, "from cached files")
    {
        set_embedder_ready(emb);
        return true;
    }
    // Don't delete files on the sync-load path — the background retry
    // loop (spawned below) will handle load failures by deleting and
    // re-downloading. The sync path is intentionally conservative
    // because a transient filesystem glitch on every embed_query/embed_document call
    // should not force a 167 MB re-download.

    // Check for tokio runtime before spawning a background task.
    // If no tokio runtime is active, we cannot spawn the background
    // download loop — mark the embedder as failed so subsequent calls
    // return false immediately.
    if tokio::runtime::Handle::try_current().is_err() {
        warn!("Embedder: no tokio runtime available — transitioning to FAILED");
        STATE.store(ModelState::Failed, Ordering::Release);
        return false;
    }
    tokio::spawn(download_retry_loop());

    false
}

// ── Public API ───────────────────────────────────────────────────────

/// Embed `text` as a query (prefixed with `"Query: "` internally).
///
/// Returns `None` if the embedder isn't available (graceful degradation).
#[must_use]
pub fn embed_query(text: &str) -> Option<Vec<f32>> {
    with_embedder(|emb| emb.embed_queries(&[text]).ok()?.into_iter().next())
}

/// Embed `text` as a document (prefixed with `"Document: "` internally).
///
/// Returns `None` if the embedder isn't available (graceful degradation).
#[must_use]
pub fn embed_document(text: &str) -> Option<Vec<f32>> {
    with_embedder(|emb| emb.embed_documents(&[text]).ok()?.into_iter().next())
}

/// Common preamble for [`embed_query`] and [`embed_document`]: ensure the
/// embedder is initialized, acquire the read lock, and pass a reference to
/// the inner [`Embedder`] to the given closure.
fn with_embedder<F>(f: F) -> Option<Vec<f32>>
where
    F: FnOnce(&Embedder) -> Option<Vec<f32>>,
{
    if !ensure_embedder() {
        return None;
    }

    let guard = GLOBAL_EMBEDDER.read().unwrap_poison();
    let emb = guard.as_ref()?;
    f(emb)
}

// ── Background download with retry ────────────────────────────────────

/// Background retry loop that downloads model and tokenizer files.
///
/// Uses exponential backoff (1 min → 2 min → 4 min → … → 30 min max) for
/// network failures. For cached-file load failures, uses a simpler strategy:
///
/// 1. If cached files exist but fail to load → delete them and re-download immediately.
/// 2. If freshly downloaded (SHA256-verified) files also fail to load → give up
///    permanently. This almost certainly indicates a code-level issue in
///    [`Embedder::load()`] that re-downloading won't fix.
///
/// # Why delete + re-download?
///
/// The root problem this solves: [`maybe_download`] returns `Ok(())` immediately if
/// the destination file already exists — it skips SHA256 re-verification on existing
/// files. If cached files become corrupted (e.g., partial write during a crash) or
/// the model format changes, the old code would loop forever: load fails →
/// maybe_download says "already have it" → load fails again → … — with no forward
/// progress.
///
/// Deleting the files on the first load failure forces a proper re-download with
/// SHA256 verification. If the freshly downloaded files also fail to load, the
/// problem is almost certainly a code-level bug, so the loop gives up and leaves
/// the global embedder uninitialized (graceful degradation to FTS-only search).
///
/// # Invariant
///
/// CONFIG must be initialized before calling this function. The only caller,
/// [`ensure_embedder`], already validates this with `.expect()` before spawning,
/// and `storage_root` is a `OnceLock` that cannot revert to `None` once set.
async fn download_retry_loop() {
    // Transitions Loading→Failed on drop if the loop is cancelled or panics.
    let _guard = ModelLoadGuard::new(&STATE);
    let models_dir = crate::util::models_dir()
        .expect("CONFIG must be initialized before embedder initialization");
    std::fs::create_dir_all(&models_dir).ok();

    let model_dest = models_dir.join(MODEL_FILENAME);
    let tokenizer_dest = models_dir.join(TOKENIZER_FILENAME);

    // Shared HTTP client reused across retries (avoids new TLS handshake per iteration).
    let Ok(client) = crate::util::http::build_download_client(MODEL_DOWNLOAD_TIMEOUT) else {
        warn!("Embedder: failed to build HTTP client — background download cancelled");
        STATE.store(ModelState::Failed, Ordering::Release);
        return;
    };

    let mut delay = Duration::from_mins(1);
    let max_delay = Duration::from_mins(30); // 30 minutes

    loop {
        // ── Phase 1: Try loading cached files ──
        // If cached files exist but fail to load, they're likely corrupted.
        // Delete them and force a fresh download with SHA256 verification.
        if model_dest.exists() && tokenizer_dest.exists() {
            if let Some(emb) = try_load_embedder(&model_dest, &tokenizer_dest, "from cached files")
            {
                set_embedder_ready(emb);
                return;
            }
            // Load failed — almost certainly file corruption. Delete and re-download
            // immediately (no backoff — we want to recover as fast as possible).
            warn!("Deleting cached model files after load failure, forcing re-download");
            let _ = std::fs::remove_file(&model_dest);
            let _ = std::fs::remove_file(&tokenizer_dest);
        }

        // ── Phase 2: Download missing files ──
        let (model_result, tokenizer_result) = tokio::join!(
            maybe_download(&client, MODEL_URL, &model_dest, Some(MODEL_SHA256)),
            maybe_download(
                &client,
                TOKENIZER_URL,
                &tokenizer_dest,
                Some(TOKENIZER_SHA256)
            ),
        );

        let model_ok = model_result.is_ok();
        let tokenizer_ok = tokenizer_result.is_ok();

        if let Err(e) = &model_result {
            warn!(error = %e, retry_after_secs = delay.as_secs(), "Failed to download embedding model, retrying");
            let _ = std::fs::remove_file(model_dest.with_extension("tmp"));
        }
        if let Err(e) = &tokenizer_result {
            warn!(error = %e, retry_after_secs = delay.as_secs(), "Failed to download tokenizer, retrying");
            let _ = std::fs::remove_file(tokenizer_dest.with_extension("tmp"));
        }

        if model_ok && tokenizer_ok {
            // Files were just downloaded and passed SHA256 verification.
            // If loading fails despite verified integrity, it's almost certainly
            // a code-level bug in Embedder::load() — re-downloading won't help.
            if let Some(emb) = try_load_embedder(&model_dest, &tokenizer_dest, "after download") {
                set_embedder_ready(emb);
                return;
            }
            warn!(
                "Giving up on embedding model: freshly downloaded, SHA256-verified files \
                 failed to load. This indicates a code-level issue with Embedder::load(). \
                 The model will remain unavailable for this session (FTS-only fallback)."
            );
            STATE.store(ModelState::Failed, Ordering::Release);
            return;
        }

        // Wait with exponential backoff before retrying a failed download.
        tokio::time::sleep(delay).await;
        delay = (delay * 2).min(max_delay);
    }
}

/// Try to load the embedder from cached files, returning [`Some`] on success.
///
/// `context` is a label included in log messages (e.g. `"from cached files"` or
/// `"after download"`) to distinguish the origin of a load attempt.
///
/// This avoids duplicating the `Embedder::load()` call + logging pattern at each site.
fn try_load_embedder(
    model_path: &Path,
    tokenizer_path: &Path,
    context: &'static str,
) -> Option<Embedder> {
    match Embedder::load(model_path, tokenizer_path) {
        Ok(emb) => {
            info!("Embedding model loaded successfully ({context})");
            Some(emb)
        }
        Err(e) => {
            warn!(reason = %e, context, "Failed to load embedding model");
            None
        }
    }
}

/// Download a file unless it already exists. Uses the shared HTTP client for
/// connection reuse across retries.
async fn maybe_download(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    expected_sha256: Option<&str>,
) -> Result<()> {
    // Skip download if the file already exists (from a previous partial success).
    if dest.exists() {
        return Ok(());
    }
    download_file(client, url, dest, expected_sha256).await
}

/// Download a single file with atomic write and size verification.
async fn download_file(
    client: &reqwest::Client,
    url: &str,
    dest: &Path,
    expected_sha256: Option<&str>,
) -> Result<()> {
    let mut size = 0u64;
    crate::util::http::download_verified(
        client,
        url,
        dest,
        expected_sha256.unwrap_or(""),
        None,
        crate::util::http::DownloadSizeCheck::Exact,
        |downloaded, _| size = downloaded,
    )
    .await?;
    info!(path = %dest.display(), size, "Downloaded model file");
    Ok(())
}

// ── GGUF metadata helpers ────────────────────────────────────────────

/// Extract a typed value from GGUF metadata.
fn get_meta<T>(
    metadata: &HashMap<String, gguf_file::Value>,
    key: &str,
    extract: impl FnOnce(&gguf_file::Value) -> std::result::Result<T, candle_core::Error>,
) -> Result<T> {
    let value = metadata
        .get(key)
        .ok_or_else(|| anyhow!("Missing metadata key '{key}'"))?;
    extract(value).map_err(|e| anyhow!("Failed to read metadata '{key}': {e}"))
}

// ── EuroBERT / LLaMA-style encoder model ────────────────────────────

/// A single transformer layer (EuroBERT = LLaMA-style encoder with SwiGLU MLP).
#[derive(Debug)]
struct Layer {
    /// Attention Q projection (no bias).
    attn_q: QMatMul,
    /// Attention K projection (no bias).
    attn_k: QMatMul,
    /// Attention V projection (no bias).
    attn_v: QMatMul,
    /// Attention output projection (no bias).
    attn_o: QMatMul,
    /// Pre-attention RMSNorm weight (1D, hidden_size).
    attn_norm: Tensor,
    /// SwiGLU gate projection (no bias).
    ffn_gate: QMatMul,
    /// SwiGLU up projection (no bias).
    ffn_up: QMatMul,
    /// SwiGLU down projection (no bias).
    ffn_down: QMatMul,
    /// Pre-FFN RMSNorm weight (1D, hidden_size).
    ffn_norm: Tensor,
}

impl Layer {
    /// Forward pass through one encoder layer.
    #[expect(clippy::many_single_char_names)]
    fn forward(
        &self,
        x: &Tensor,
        mask: &Tensor,
        cos: &Tensor,
        sin: &Tensor,
        n_head: usize,
        head_dim: usize,
    ) -> Result<Tensor> {
        // ── Self-attention with pre-norm ──
        let residual = x;
        let h = candle_nn::ops::rms_norm(x, &self.attn_norm, 1e-5)?;

        // Project to Q, K, V
        let q = self.attn_q.forward(&h)?;
        let k = self.attn_k.forward(&h)?;
        let v = self.attn_v.forward(&h)?;

        // Multi-head attention
        let h = Layer::apply_attention(&q, &k, &v, mask, cos, sin, n_head, head_dim)?;
        let h = self.attn_o.forward(&h)?;
        let h = (h + residual)?;

        // ── SwiGLU MLP with pre-norm ──
        let residual = &h;
        let h = candle_nn::ops::rms_norm(&h, &self.ffn_norm, 1e-5)?;

        let gate = self.ffn_gate.forward(&h)?;
        let up = self.ffn_up.forward(&h)?;
        let h = self
            .ffn_down
            .forward(&(candle_nn::ops::silu(&gate)? * up)?)?;
        let h = (h + residual)?;

        Ok(h)
    }

    /// Bidirectional multi-head attention with RoPE (no KV cache).
    #[expect(clippy::too_many_arguments)]
    fn apply_attention(
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        mask: &Tensor,
        cos: &Tensor,
        sin: &Tensor,
        n_head: usize,
        head_dim: usize,
    ) -> Result<Tensor> {
        let (b_sz, seq_len, n_embd) = q.shape().dims3()?;

        // Reshape: [batch, seq, n_head * head_dim] -> [batch, seq, n_head, head_dim] -> [batch, n_head, seq, head_dim]
        let q = q
            .reshape((b_sz, seq_len, n_head, head_dim))?
            .transpose(1, 2)?;
        let k = k
            .reshape((b_sz, seq_len, n_head, head_dim))?
            .transpose(1, 2)?;
        let v = v
            .reshape((b_sz, seq_len, n_head, head_dim))?
            .transpose(1, 2)?
            .contiguous()?;

        // Apply RoPE via rope_slow.
        // cos/sin have shape [MAX_SEQ_LEN, head_dim/2] as produced by
        // precompute_freqs_cis; rope_slow internally duplicates them along
        // the last dimension to match head_dim.
        let q = rope_slow(&q, cos, sin)?;
        let k = rope_slow(&k, cos, sin)?;

        // Scaled dot-product attention (no causal mask — full bidirectional)
        #[expect(clippy::cast_precision_loss)]
        let scale = 1.0_f64 / (head_dim as f64).sqrt();
        let att = q.matmul(&k.t()?)?;
        let att = (att * scale)?;
        let mask = mask.broadcast_as(att.shape())?;
        let att = (att + mask)?;
        let att = candle_nn::ops::softmax_last_dim(&att)?;
        let y = att.matmul(&v)?;

        // Reshape back: [batch, n_head, seq, head_dim] -> [batch, seq, n_embd]
        let y = y.transpose(1, 2)?.reshape((b_sz, seq_len, n_embd))?;
        Ok(y)
    }
}

// ── Weight-loading helpers for `Embedder::load` ─────────────────────

/// Load a quantized matrix-multiply weight tensor from a GGUF file.
fn load_qmatmul(
    content: &gguf_file::Content,
    file: &mut File,
    prefix: &str,
    name: &str,
    device: &Device,
) -> Result<QMatMul> {
    QMatMul::from_qtensor(
        content
            .tensor(file, &format!("{prefix}.{name}.weight"), device)
            .with_context(|| format!("Failed to load {prefix}.{name}.weight"))?,
    )
    .with_context(|| format!("Failed to create QMatMul for {name}"))
}

/// Load and dequantize a norm weight tensor from a GGUF file.
fn load_norm(
    content: &gguf_file::Content,
    file: &mut File,
    prefix: &str,
    name: &str,
    device: &Device,
) -> Result<Tensor> {
    content
        .tensor(file, &format!("{prefix}.{name}.weight"), device)
        .with_context(|| format!("Failed to load {prefix}.{name}.weight"))?
        .dequantize(device)
        .with_context(|| format!("Failed to dequantize {name}"))
}

// ── Embedder ─────────────────────────────────────────────────────────

/// The embedding model: EuroBERT encoder + tokenizer + pooling.
pub struct Embedder {
    tokenizer: Tokenizer,
    tok_embeddings: Embedding,
    layers: Vec<Layer>,
    output_norm: Tensor,
    cos: Tensor,
    sin: Tensor,
    head_dim: usize,
    n_head: usize,
    pad_id: u32,
}

impl Embedder {
    /// Load an [`Embedder`] from cached GGUF and tokenizer files.
    ///
    /// Does NOT download — the caller is responsible for ensuring the files exist.
    /// Returns an error if files are missing, corrupted, or the model architecture
    /// is unexpected.
    #[expect(clippy::too_many_lines)]
    pub fn load(model_path: &Path, tokenizer_path: &Path) -> Result<Self> {
        let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(|e| {
            anyhow!(
                "Failed to load tokenizer from {}: {e}",
                tokenizer_path.display()
            )
        })?;

        // Discover pad token ID from the tokenizer.
        // jina-embeddings-v5 uses eos_token_id = 128001 as pad token.
        let pad_id = tokenizer
            .token_to_id("<|end_of_text|>")
            .or_else(|| tokenizer.token_to_id("<|pad|>"))
            .or_else(|| tokenizer.token_to_id("[PAD]"))
            .unwrap_or(DEFAULT_PAD_ID);

        debug!(pad_id, "Discovered pad token ID from tokenizer");

        let device = Device::Cpu;

        // Open and read GGUF file
        let mut file = std::fs::File::open(model_path)
            .map_err(|e| anyhow!("Failed to open model file {}: {e}", model_path.display()))?;
        let content = gguf_file::Content::read(&mut file)
            .map_err(|e| anyhow!("Failed to read GGUF file: {e}"))?;

        // Read architecture metadata
        let hidden_size = get_meta(&content.metadata, "eurobert.embedding_length", |v| {
            v.to_u32()
        })? as usize;
        let n_head = get_meta(&content.metadata, "eurobert.attention.head_count", |v| {
            v.to_u32()
        })? as usize;
        let head_dim = get_meta(&content.metadata, "eurobert.attention.value_length", |v| {
            v.to_u32()
        })? as usize;
        let rope_freq_base = get_meta(
            &content.metadata,
            "eurobert.rope.freq_base",
            candle_core::quantized::gguf_file::Value::to_f32,
        )
        .unwrap_or(ROPE_FREQ_BASE);

        // Count layers by scanning tensor names
        let n_layers = content
            .tensor_infos
            .keys()
            .filter_map(|name| {
                let name = name.as_str();
                if name.starts_with("blk.") && name.ends_with(".attn_q.weight") {
                    // Extract layer index
                    name.trim_start_matches("blk.")
                        .split('.')
                        .next()?
                        .parse::<usize>()
                        .ok()
                } else {
                    None
                }
            })
            .max()
            .map(|max| max + 1)
            .context("No layer tensors found in GGUF file")?;

        info!(
            hidden_size,
            n_head, head_dim, n_layers, rope_freq_base, "Loading EuroBERT embedding model"
        );

        // ── Load token embeddings (dequantize for use with candle_nn::Embedding) ──
        let tok_embd_qt = content
            .tensor(&mut file, "token_embd.weight", &device)
            .context("Failed to load token_embd.weight")?;
        let tok_embd_f32 = tok_embd_qt
            .dequantize(&device)
            .context("Failed to dequantize token_embd.weight")?;
        let tok_embeddings = Embedding::new(tok_embd_f32, hidden_size);

        // ── Load output norm ──
        let output_norm_qt = content
            .tensor(&mut file, "output_norm.weight", &device)
            .context("Failed to load output_norm.weight")?;
        let output_norm = output_norm_qt
            .dequantize(&device)
            .context("Failed to dequantize output_norm.weight")?;

        // ── Load transformer layers ──
        let mut layers = Vec::with_capacity(n_layers);
        for i in 0..n_layers {
            let prefix = format!("blk.{i}");

            let attn_q = load_qmatmul(&content, &mut file, &prefix, "attn_q", &device)?;
            let attn_k = load_qmatmul(&content, &mut file, &prefix, "attn_k", &device)?;
            let attn_v = load_qmatmul(&content, &mut file, &prefix, "attn_v", &device)?;
            let attn_o = load_qmatmul(&content, &mut file, &prefix, "attn_output", &device)?;
            let attn_norm = load_norm(&content, &mut file, &prefix, "attn_norm", &device)?;
            let ffn_gate = load_qmatmul(&content, &mut file, &prefix, "ffn_gate", &device)?;
            let ffn_up = load_qmatmul(&content, &mut file, &prefix, "ffn_up", &device)?;
            let ffn_down = load_qmatmul(&content, &mut file, &prefix, "ffn_down", &device)?;
            let ffn_norm = load_norm(&content, &mut file, &prefix, "ffn_norm", &device)?;

            layers.push(Layer {
                attn_q,
                attn_k,
                attn_v,
                attn_o,
                attn_norm,
                ffn_gate,
                ffn_up,
                ffn_down,
                ffn_norm,
            });
        }

        // ── Precompute RoPE frequencies ──
        let (cos, sin) = precompute_freqs_cis(head_dim, rope_freq_base, &device)?;

        // ── Build embedder ──
        let emb = Self {
            tokenizer,
            tok_embeddings,
            layers,
            output_norm,
            cos,
            sin,
            head_dim,
            n_head,
            pad_id,
        };

        // ── Warm-up: run a single short input to validate the model ──
        let v = emb.embed_documents(&["."])?;
        anyhow::ensure!(
            !v.is_empty() && !v[0].is_empty(),
            "Embedder warm-up produced empty output"
        );
        anyhow::ensure!(
            v[0].len() == hidden_size,
            "Embedder warm-up produced wrong dimension: expected {hidden_size}, got {}",
            v[0].len()
        );

        info!("Embedder initialized successfully (hidden_size={hidden_size}, layers={n_layers})");
        Ok(emb)
    }

    /// Embed texts as queries (prefixed with `"Query: "`).
    pub fn embed_queries(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        self.embed_prefixed("Query: ", texts)
    }

    /// Embed texts as documents (prefixed with `"Document: "`).
    pub fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        self.embed_prefixed("Document: ", texts)
    }

    /// Core embedding method.
    fn embed_prefixed(&self, prefix: &str, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        // ── Tokenize ──
        let encodings: Vec<_> = texts
            .iter()
            .map(|t| {
                let input = format!("{prefix}{t}");
                self.tokenizer.encode(input, true)
            })
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(|e| anyhow!("Tokenization error: {e}"))?;

        if encodings.is_empty() {
            anyhow::bail!("Empty input");
        }

        // Determine max sequence length across batch (clamped to MAX_SEQ_LEN)
        let max_len = encodings
            .iter()
            .map(|e| e.len().min(MAX_SEQ_LEN))
            .max()
            .context("Empty encoding")?;

        let batch_size = encodings.len();

        // ── Build input_ids and attention_mask ──
        let mut input_ids_vec = vec![i64::from(self.pad_id); batch_size * max_len];
        let mut attention_mask_vec = vec![0i64; batch_size * max_len];

        for (row, enc) in encodings.iter().enumerate() {
            let ids = enc.get_ids();
            let mask = enc.get_attention_mask();
            let len = ids.len().min(MAX_SEQ_LEN);
            for col in 0..len {
                input_ids_vec[row * max_len + col] = i64::from(ids[col]);
                attention_mask_vec[row * max_len + col] = i64::from(mask[col]);
            }
        }

        let input_ids = Tensor::from_vec(input_ids_vec, (batch_size, max_len), &Device::Cpu)?;
        let attention_mask =
            Tensor::from_vec(attention_mask_vec, (batch_size, max_len), &Device::Cpu)?;

        // ── Forward pass through the model ──
        let embeddings = self.forward(&input_ids, &attention_mask)?;

        // ── Post-process: last-token pooling + L2 normalization ──
        let result = last_token_pool_and_l2_normalize(&embeddings, &attention_mask)?;

        Ok(result)
    }

    /// Full model forward pass.
    fn forward(&self, input_ids: &Tensor, attention_mask: &Tensor) -> Result<Tensor> {
        let (_batch_size, _seq_len) = input_ids.shape().dims2()?;

        // No causal masking — this is an encoder
        let mask = build_attn_mask(attention_mask, &Device::Cpu)?;

        let mut h = self.tok_embeddings.forward(input_ids)?;
        // h: [batch, seq, hidden_size]

        for layer in &self.layers {
            h = layer.forward(&h, &mask, &self.cos, &self.sin, self.n_head, self.head_dim)?;
        }

        h = candle_nn::ops::rms_norm(&h, &self.output_norm, 1e-5)?;

        Ok(h)
    }
}

// ── RoPE ─────────────────────────────────────────────────────────────

/// Precompute cosine and sine tables for rotary position embeddings.
///
/// Returns `(cos, sin)` each with shape `[MAX_SEQ_LEN, head_dim/2]`.
/// These are consumed by `rope_slow` which internally duplicates them
/// along the last dimension to match `head_dim`.
fn precompute_freqs_cis(
    head_dim: usize,
    freq_base: f32,
    device: &Device,
) -> Result<(Tensor, Tensor)> {
    #[expect(clippy::cast_precision_loss)]
    let theta: Vec<f32> = (0..head_dim)
        .step_by(2)
        .map(|i| 1.0_f32 / freq_base.powf(i as f32 / head_dim as f32))
        .collect();

    let theta = Tensor::from_vec(theta, (head_dim / 2,), device)?;
    #[expect(clippy::cast_possible_truncation)]
    let positions = Tensor::arange(0u32, MAX_SEQ_LEN as u32, device)?
        .to_dtype(DType::F32)?
        .reshape((MAX_SEQ_LEN, 1))?;

    let idx_theta = positions.matmul(&theta.reshape((1, theta.elem_count()))?)?;
    let cos = idx_theta.cos()?;
    let sin = idx_theta.sin()?;

    Ok((cos, sin))
}

// ── Attention mask ───────────────────────────────────────────────────

/// Build a bidirectional attention mask from a tokenizer attention mask.
///
/// The input `attention_mask` has shape `[batch, seq]` with 1 for real tokens and
/// 0 for padding. The output has shape `[batch, 1, seq, seq]` where:
/// - Entry (i,j) is 0 if both positions i and j are real tokens,
/// - Otherwise it's a large negative value (-1e10) that acts as -inf for softmax.
///
/// Uses `-1e10` instead of `f32::NEG_INFINITY` because `0 * NEG_INFINITY = NaN`
/// per IEEE 754, which would corrupt the entire attention computation.
fn build_attn_mask(attention_mask: &Tensor, device: &Device) -> Result<Tensor> {
    let (batch_size, seq_len) = attention_mask.shape().dims2()?;

    // Expand to [batch, 1, seq, seq]: mask[i, j] = mask[i] * mask[j]
    // (both tokens must be real to attend)
    let mask_f32 = attention_mask.to_dtype(DType::F32)?;
    let mask_a = mask_f32.reshape((batch_size, 1, 1, seq_len))?;
    let mask_b = mask_f32.reshape((batch_size, 1, seq_len, 1))?;
    let pairwise = mask_a.broadcast_mul(&mask_b)?;
    // pairwise now has 1 where both are real, 0 where either is padding

    // Convert to attention mask using where_cond:
    // - Where pairwise == 1 (attend): mask = 0
    // - Where pairwise == 0 (masked): mask = -1e10
    // Using -1e10 instead of NEG_INFINITY because 0 * NEG_INFINITY = NaN.
    let large_neg = Tensor::new(-1e10_f32, device)?.broadcast_as(pairwise.shape())?;
    let zero = Tensor::new(0.0_f32, device)?.broadcast_as(pairwise.shape())?;
    // Build boolean predicate: pairwise == 0
    let mask_cond = pairwise.eq(&zero)?;
    let mask = mask_cond.where_cond(&large_neg, &zero)?;

    Ok(mask)
}

// ── Pooling and normalization ────────────────────────────────────────

/// Extract embeddings via last-token pooling and L2 normalize.
///
/// Takes the embedding at the last non-padding token position for each sequence,
/// then L2-normalizes each vector.
fn last_token_pool_and_l2_normalize(
    embeddings: &Tensor,
    attention_mask: &Tensor,
) -> Result<Vec<Vec<f32>>> {
    let (batch_size, _seq_len, hidden_size) = embeddings.shape().dims3()?;

    let mut results = Vec::with_capacity(batch_size);

    // Sum attention mask along seq dimension to find last real token position
    // last_pos = sum(mask) - 1 (0-indexed)
    let seq_lengths: Vec<i64> = attention_mask.sum(1)?.to_vec1()?;

    for (i, &seq_len) in seq_lengths.iter().enumerate().take(batch_size) {
        #[expect(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
        let last_pos = (seq_len - 1).max(0) as usize;

        // Extract embedding at last position: [batch, seq, hidden] -> [hidden]
        let token_emb = embeddings.narrow(0, i, 1)?.narrow(1, last_pos, 1)?;
        let token_emb = token_emb.reshape(hidden_size)?;

        // L2 normalize
        let norm = token_emb
            .sqr()?
            .sum_all()?
            .sqrt()?
            .to_scalar::<f32>()?
            .max(1e-12);
        let normalized = token_emb.broadcast_div(&Tensor::new(norm, token_emb.device())?)?;

        let vec: Vec<f32> = normalized.to_vec1()?;
        results.push(vec);
    }

    Ok(results)
}

// ── tests ─────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::vector::cosine_similarity;
    use std::sync::Mutex;

    /// Global mutex to serialize embedder tests that manipulate the
    /// global [`STATE`] machine. Without it, parallel test execution
    /// can race on `reset_global_state()`, `ensure_embedder()`, and
    /// direct `STATE` stores, producing non-deterministic failures.
    static TEST_GLOBAL_STATE_MUTEX: Mutex<()> = Mutex::new(());

    /// Initialize config storage root for tests using the shared test root.
    /// Returns the root path (used as storage root).
    ///
    /// Uses the single process-level test root (not a per-test temp dir), so
    /// there is exactly one test-owned temp root per process, cleaned up at
    /// process exit.
    fn init_test_config() -> std::path::PathBuf {
        let root = crate::util::test::test_root().clone();
        let _ = crate::config::CONFIG.try_set_storage_root(root.clone());
        root
    }

    /// Like [`test_embedder_opt`] but panics with a clear message if model
    /// files are not cached on disk.
    fn test_embedder_opt() -> Option<&'static Embedder> {
        use std::sync::OnceLock;

        // Share a single model load across all tests via OnceLock.
        // The first test to call this pays the load (~2-3 s for the
        // 157 MB GGUF on this machine); subsequent tests reuse the
        // already-loaded embedder.
        static TEST_EMBEDDER: OnceLock<Option<Embedder>> = OnceLock::new();

        TEST_EMBEDDER
            .get_or_init(|| {
                // Collect all candidate models directories (deduplicated).
                let mut candidates = Vec::new();

                // 1. CONFIG storage root (the shared test root in test runs —
                //    empty models/ dir unless the ignored TTS e2e test's
                //    symlink is present in an --include-ignored run).
                if let Some(root) = crate::config::CONFIG.try_storage_root() {
                    candidates.push(root.join("models"));
                }

                // 2. Real home directory cache (always present in dev/CI environments).
                if let Some(home) = std::env::var("HOME").ok().filter(|h| !h.is_empty()) {
                    let real = std::path::PathBuf::from(&home)
                        .join(".mahbot")
                        .join("models");
                    if !candidates.contains(&real) {
                        candidates.push(real);
                    }
                }

                // Try each candidate until we find model files.
                for models_dir in &candidates {
                    let model_path = models_dir.join(MODEL_FILENAME);
                    let tokenizer_path = models_dir.join(TOKENIZER_FILENAME);

                    if model_path.exists() && tokenizer_path.exists() {
                        match Embedder::load(&model_path, &tokenizer_path) {
                            Ok(emb) => return Some(emb),
                            Err(e) => {
                                eprintln!("WARNING: Failed to load test embedder: {e}");
                                return None;
                            }
                        }
                    }
                }

                // No model files found in any candidate directory.
                let last_candidate = candidates.last().map(|p| p.display().to_string());
                eprintln!(
                    "Embedder model files not found. Looked in: {}. \
                     Run the application first to download embedding models (~150 MB).",
                    last_candidate.as_deref().unwrap_or("<none>")
                );
                None
            })
            .as_ref()
    }

    /// Reset global embedder state for hermetic testing.
    fn reset_global_state() {
        *GLOBAL_EMBEDDER.write().unwrap_poison() = None;
        STATE.store(ModelState::Uninit, Ordering::Release);
    }

    #[tokio::test]
    async fn test_embedder_graceful_degradation() {
        let _lock = TEST_GLOBAL_STATE_MUTEX.lock().unwrap();
        // Use a temp dir as storage root — no model files there.
        let _root = init_test_config();
        reset_global_state();

        // verify: without model files, embed_document() returns None
        let result = embed_document("test");
        assert!(
            result.is_none(),
            "embed_document() should return None when model not available"
        );

        // Verify the global embedder is still empty
        let guard = GLOBAL_EMBEDDER.read().unwrap_poison();
        assert!(guard.is_none(), "global embedder should remain None");
    }

    #[test]
    fn test_ensure_embedder_terminally_failed() {
        let _lock = TEST_GLOBAL_STATE_MUTEX.lock().unwrap();
        // Regression test: when STATE is FAILED, ensure_embedder() should return
        // false immediately (without attempting re-initialization).
        let _root = init_test_config();
        reset_global_state();

        // Manually set STATE to FAILED (simulates the terminal condition from
        // download_retry_loop after freshly-downloaded files fail to load).
        STATE.store(ModelState::Failed, Ordering::Release);

        // First call: should return false immediately
        let r1 = ensure_embedder();
        assert!(
            !r1,
            "ensure_embedder should return false in ModelState::Failed"
        );
        assert_eq!(
            STATE.load(Ordering::Acquire),
            ModelState::Failed,
            "STATE should remain FAILED after ensure_embedder()"
        );

        // Second call: should again return false (no state change)
        let r2 = ensure_embedder();
        assert!(!r2, "ensure_embedder should still return false");
        assert_eq!(
            STATE.load(Ordering::Acquire),
            ModelState::Failed,
            "STATE should remain FAILED after second call"
        );
        // Clean up state for subsequent tests
        reset_global_state();
    }

    /// Test that [`ensure_embedder()`] handles the absence of a tokio runtime
    /// gracefully by transitioning [`STATE`] to [`ModelState::Failed`].
    ///
    /// This is a regular `#[test]`, not `#[tokio::test]`, so there is no tokio
    /// runtime active. When `ensure_embedder()` reaches the runtime check via
    /// [`tokio::runtime::Handle::try_current()`], it detects the absence and sets
    /// [`ModelState::Failed`], preventing a permanent wedge.
    #[test]
    fn test_ensure_embedder_no_runtime() {
        let _lock = TEST_GLOBAL_STATE_MUTEX.lock().unwrap();
        let _root = init_test_config();
        reset_global_state();

        // ensure_embedder() will CAS → LOADING, try the sync load (fails
        // because there are no model files in the temp dir), then check
        // Handle::try_current() — which fails without a tokio runtime.
        let result = ensure_embedder();

        assert!(
            !result,
            "ensure_embedder should return false when no tokio runtime is available"
        );
        assert_eq!(
            STATE.load(Ordering::Acquire),
            ModelState::Failed,
            "STATE should be FAILED after detecting no tokio runtime"
        );

        // Subsequent calls immediately return false (no re-attempt).
        let r2 = ensure_embedder();
        assert!(!r2, "ensure_embedder should still return false");
        assert_eq!(
            STATE.load(Ordering::Acquire),
            ModelState::Failed,
            "STATE should remain FAILED"
        );

        reset_global_state();
    }

    #[test]
    fn test_load_corrupted_files_fails() {
        // Create corrupted model and tokenizer files in a temp directory
        // and verify that Embedder::load() correctly reports failure.
        let tmp = tempfile::TempDir::new().expect("failed to create temp dir");
        let model_path = tmp.path().join(MODEL_FILENAME);
        let tokenizer_path = tmp.path().join(TOKENIZER_FILENAME);

        // Write invalid data: model file gets random bytes, tokenizer gets
        // non-JSON text (Tokenizer::from_file will fail early).
        std::fs::write(&model_path, b"not a valid gguf file").unwrap();
        std::fs::write(&tokenizer_path, b"not valid json at all").unwrap();

        let result = Embedder::load(&model_path, &tokenizer_path);
        assert!(
            result.is_err(),
            "Embedder::load() should fail on corrupted files"
        );

        // Verify that the error includes one of the file paths so the user can
        // locate the problematic file. The assertion checks either path because
        // the loading order (tokenizer first vs model first) is an implementation
        // detail that should not make the test brittle.
        let err = result.err().expect("just checked is_err");
        let err_msg = format!("{err:#}");
        let model_path_str = model_path.to_string_lossy();
        let tokenizer_path_str = tokenizer_path.to_string_lossy();
        assert!(
            err_msg.contains(model_path_str.as_ref())
                || err_msg.contains(tokenizer_path_str.as_ref()),
            "Error should mention one of the file paths, got: {err_msg}"
        );
    }

    /// Embeddings from the real cached jina-embeddings-v5 GGUF model are
    /// 768-dimensional L2-normalized unit vectors — for single and batch
    /// document inputs and for query inputs.
    ///
    /// This test is `#[ignore]` by default because it loads the real cached
    /// 157 MB GGUF model (~3-4 s). Run it explicitly with:
    ///
    /// ```sh
    /// cargo test test_embedding_produces_unit_vectors -- --ignored --nocapture
    /// ```
    #[ignore = "loads the real cached 157 MB GGUF model (~3-4 s); runs only when explicitly invoked"]
    #[test]
    fn test_embedding_produces_unit_vectors() {
        let Some(emb) = test_embedder_opt() else {
            eprintln!("SKIP: embedder model files not cached — skipping model-backed test");
            return;
        };

        // embed_documents with single input
        let v = emb.embed_documents(&["hello world"]).unwrap();
        assert_eq!(v.len(), 1);
        // jina-embeddings-v5 produces 768-dimensional vectors
        assert_eq!(v[0].len(), 768);
        // L2-normalized → unit vector (approximately norm 1)
        let norm: f32 = v[0].iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (norm - 1.0).abs() < 1e-5,
            "expected unit vector, got norm={norm}"
        );

        // embed_documents with multiple inputs
        let docs = &["first document", "second document about something"];
        let v = emb.embed_documents(docs).unwrap();
        assert_eq!(v.len(), 2);
        for vec in &v {
            assert_eq!(vec.len(), 768);
            let norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
            assert!(
                (norm - 1.0).abs() < 1e-5,
                "expected unit vector, got norm={norm}"
            );
        }

        // embed_queries
        let v = emb.embed_queries(&["what is rust?"]).unwrap();
        assert_eq!(v.len(), 1);
        assert_eq!(v[0].len(), 768);
        let norm: f32 = v[0].iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (norm - 1.0).abs() < 1e-5,
            "expected unit vector, got norm={norm}"
        );
    }

    /// Cosine similarity ranks same-topic embeddings above cross-topic ones:
    /// two rust documents must be more similar than a rust and a python
    /// document.
    ///
    /// This test is `#[ignore]` by default because it loads the real cached
    /// 157 MB GGUF model (~3-4 s). Run it explicitly with:
    ///
    /// ```sh
    /// cargo test test_similar_embeddings_are_similar -- --ignored --nocapture
    /// ```
    #[ignore = "loads the real cached 157 MB GGUF model (~3-4 s); runs only when explicitly invoked"]
    #[test]
    fn test_similar_embeddings_are_similar() {
        let Some(emb) = test_embedder_opt() else {
            eprintln!("SKIP: embedder model files not cached — skipping model-backed test");
            return;
        };
        let v = emb
            .embed_documents(&[
                "rust programming language",
                "the rust programming language",
                "python programming language",
            ])
            .unwrap();
        let sim_01 = cosine_similarity(&v[0], &v[1]);
        let sim_02 = cosine_similarity(&v[0], &v[2]);
        assert!(
            sim_01 > sim_02,
            "rust/rust ({sim_01}) should be more similar than rust/python ({sim_02})"
        );
    }

    /// Embedding an empty document list must produce an error, not a panic or
    /// an empty result.
    ///
    /// This test is `#[ignore]` by default because it loads the real cached
    /// 157 MB GGUF model (~3-4 s). Run it explicitly with:
    ///
    /// ```sh
    /// cargo test test_empty_input_fails -- --ignored --nocapture
    /// ```
    #[ignore = "loads the real cached 157 MB GGUF model (~3-4 s); runs only when explicitly invoked"]
    #[test]
    fn test_empty_input_fails() {
        let Some(emb) = test_embedder_opt() else {
            eprintln!("SKIP: embedder model files not cached — skipping model-backed test");
            return;
        };
        let result = emb.embed_documents(&[]);
        assert!(result.is_err(), "empty input should produce an error");
    }
}