//! MiniLM-L6-v2 embedder via fastembed (bundled ONNX Runtime).
//!
//! Default and currently only embedder. Output is 384-dim float32.

use anyhow::{Context, Result};
use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};

use super::Embedder;

/// Stable identifier for this embedder. Persisted in the manifest and used
/// for mismatch detection at search time.
pub const MINILM_ID: &str = "minilm-l6-v2";

/// Embedding dimension (384) for all-MiniLM-L6-v2.
pub const MINILM_DIM: u32 = 384;

/// Character budget for the context string fed to MiniLM-L6-v2.
/// MiniLM accepts ~256 wordpiece tokens (~4.5 chars/token → ~1100 chars).
pub const MINILM_CHAR_BUDGET: usize = 1100;

/// Miss-count threshold below which `Device::Auto` stays on CPU for MiniLM —
/// the GPU warm-up isn't worth it for a tiny `vex update`. MiniLM (~22M) has a
/// low per-symbol CPU cost, so its GPU break-even is high. Measured on an RTX
/// 3080 / CommonLib (27,997 symbols, docs/GPU_SUPPORT.md §11): MiniLM CPU embed
/// ≈ 59 sym/s vs CUDA ≈ 3043 sym/s with a ~10 s model-load warm-up, putting the
/// CPU-vs-GPU break-even at ≈ 600 misses; 512 sits just under that. (Above the
/// gate, GPU is a large win even for MiniLM — 51× CUDA / 29× DirectML on embed —
/// so this gate only protects small incremental updates, not full indexes.)
/// Heavier models override this with smaller values (see
/// `crate::embed::extra::Spec` — crate-internal, so no doc-link).
pub const MINILM_GPU_AUTO_MIN_MISSES: usize = 512;

pub struct MiniLMEmbedder {
    model: TextEmbedding,
    /// True when a GPU execution provider was registered, so `embed_batch`
    /// uses length-aware micro-batching (bounds VRAM + avoids padding waste).
    gpu: bool,
}

impl MiniLMEmbedder {
    /// Initialize the model. Downloads ~86 MB on first use into the
    /// global vex embedding cache (shared across projects). Without an
    /// explicit cache dir fastembed would drop `.fastembed_cache/` into
    /// the current working directory — that re-downloads the same model
    /// for every project and pollutes the project tree.
    /// Construct on the default CPU device. Equivalent to
    /// `with_device(Device::Cpu, false)` — the library floor. The GPU decision
    /// lives in the index path (see [`crate::embed::make_embedder_with_device`]).
    pub fn new() -> Result<Self> {
        Self::with_device(crate::embed::device::Device::Cpu, false)
    }

    /// Construct on a specific compute [`Device`](crate::embed::device::Device).
    /// `Device::Cpu` (and `Auto` on a CPU-only build) yields an empty
    /// execution-provider list — byte-for-byte the legacy CPU load. GPU
    /// providers are chained via fastembed's `with_execution_providers`; the
    /// downloaded ONNX bytes are identical regardless of EP, so the integrity
    /// check below is unaffected. `strict` makes a failed EP registration a
    /// hard error instead of a silent CPU fallback (the `vex gpu` probe).
    pub fn with_device(device: crate::embed::device::Device, strict: bool) -> Result<Self> {
        let cache_dir = crate::util::config::embed_cache_dir();
        // Surface the cache-dir creation error explicitly. Without
        // context, a failure here would propagate up as the cryptic
        // "failed to load MiniLM-L6-v2 embedding model" generated by
        // fastembed when it cannot read the (still-missing) directory.
        std::fs::create_dir_all(&cache_dir).with_context(|| {
            format!(
                "create embedding cache directory at {}",
                cache_dir.display()
            )
        })?;
        let execution_providers = crate::embed::device::execution_providers(device, strict)?;
        let gpu = !execution_providers.is_empty();
        let model = TextEmbedding::try_new(
            InitOptions::new(EmbeddingModel::AllMiniLML6V2)
                .with_cache_dir(cache_dir.clone())
                .with_show_download_progress(true)
                .with_execution_providers(execution_providers),
        )
        .context("failed to load MiniLM-L6-v2 embedding model")?;

        // Integrity check (10.4): verify the downloaded ONNX bytes match
        // the pinned SHA-256. Best-effort defence-in-depth — the model
        // is already loaded into ONNX Runtime by the time we run, but
        // a mismatch immediately tells the user the cache or CDN gave
        // them an unexpected file and the resulting index should not
        // be trusted. `find_minilm_onnx` returns `None` when fastembed
        // cache layout has drifted; in that case we warn and proceed
        // so a fastembed upgrade doesn't break vex for everyone.
        match super::integrity::find_minilm_onnx(&cache_dir) {
            Some(onnx_path) => {
                // v1.13 P2: `verify_with_marker` caches `(mtime, size,
                // sha)` next to the ONNX so the warm path skips the
                // ~163 ms full hash. First run hashes + writes marker;
                // every subsequent `vex search --semantic` is ~μs.
                super::integrity::verify_with_marker(
                    &onnx_path,
                    super::integrity::MINILM_ONNX_SHA256,
                )?;
            }
            None => {
                tracing::warn!(
                    cache_dir = %cache_dir.display(),
                    "MiniLM ONNX file not found at the expected HF-Hub-style \
                     path; skipping integrity check (fastembed cache layout \
                     may have changed — please report)"
                );
            }
        }

        Ok(Self { model, gpu })
    }
}

impl Embedder for MiniLMEmbedder {
    fn id(&self) -> &'static str {
        MINILM_ID
    }

    fn dim(&self) -> u32 {
        MINILM_DIM
    }

    fn char_budget(&self) -> usize {
        MINILM_CHAR_BUDGET
    }

    fn embed(&mut self, text: &str) -> Result<Vec<f32>> {
        let results = self.model.embed(vec![text], None)?;
        results
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("embedding model returned no vector for input"))
    }

    fn embed_batch(&mut self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        if self.gpu {
            // GPU: length-aware micro-batching bounds VRAM and avoids padding
            // waste (see crate::embed::batching).
            return crate::embed::batching::embed_length_aware(&mut self.model, texts);
        }
        // CPU: fastembed's internal batching (default 256) — unchanged.
        let refs: Vec<&str> = texts.iter().map(String::as_str).collect();
        let results = self.model.embed(refs, None)?;
        Ok(results)
    }
}
