embedd 0.4.0

//! `embedd`: embedding interfaces + reusable backends (multi-modality substrate).
//!
//! This crate is the "shared embedding substrate": consumers should depend on `embedd`
//! (traits + basic types), and enable backend **features** (`candle-hf`, `openai`, `tei`, etc.)
//! as needed.
//!
//! Scope: **any modality** (text, images, audio, …) as long as the interface remains small and
//! testable. Concretely, we provide:
//! - `TextEmbedder` (text -> vectors)
//! - `ImageEmbedder` (bytes -> vectors)
//! - `AudioEmbedder` (bytes -> vectors; placeholder contract)
//! - extension traits for token-level / sparse embeddings.
//!
//! ## Environment variables
//!
//! Primary (`EMBEDD_*`):
//! - `EMBEDD_MODEL` / `EMBEDD_MODEL_DIR` -- model source
//! - `EMBEDD_MAX_LEN` -- max token length
//! - `EMBEDD_QUERY_PREFIX` / `EMBEDD_DOC_PREFIX` -- prompt prefixes
//!
//! Legacy `IKSH_*` equivalents are supported as fallback via `from_env_any()` but deprecated.

// ---------------------------------------------------------------------------
// Qdrant integration (feature-gated)
// ---------------------------------------------------------------------------

#[cfg(feature = "qdrant")]
pub mod qdrant;

/// Whether an embedding is for a query or a document/passage.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum EmbedMode {
    /// Query embedding (may use a different instruction/prefix).
    Query,
    /// Document / passage embedding.
    Document,
}

/// Prompt template applied before tokenization, used for instruction-tuned / prompt-tuned embedders.
///
/// Common patterns in the ecosystem:
/// - **E5**: `query: {text}` vs `passage: {text}`
/// - **BGE** (often): no prompt, or `query: ...` / `passage: ...` depending on checkpoint
/// - **Instruct models**: sometimes use longer task prompts; we keep this as plain prefixing
///   for now because it composes with both local and remote backends.
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct PromptTemplate {
    pub query_prefix: String,
    pub doc_prefix: String,
}

impl Default for PromptTemplate {
    fn default() -> Self {
        Self {
            query_prefix: "query: ".to_string(),
            doc_prefix: "passage: ".to_string(),
        }
    }
}

impl PromptTemplate {
    pub fn apply(&self, mode: EmbedMode, text: &str) -> String {
        match mode {
            EmbedMode::Query => format!("{}{}", self.query_prefix, text),
            EmbedMode::Document => format!("{}{}", self.doc_prefix, text),
        }
    }

    /// Load prompt prefixes from the legacy `iksh` env vars (compat).
    #[deprecated(since = "0.2.0", note = "use from_embedd_env() or from_env_any()")]
    pub fn from_iksh_env() -> Self {
        let query_prefix = std::env::var("IKSH_EMBED_QUERY_PREFIX")
            .unwrap_or_else(|_| Self::default().query_prefix);
        let doc_prefix =
            std::env::var("IKSH_EMBED_DOC_PREFIX").unwrap_or_else(|_| Self::default().doc_prefix);
        Self {
            query_prefix,
            doc_prefix,
        }
    }

    /// Load prompt prefixes from `EMBEDD_QUERY_PREFIX` / `EMBEDD_DOC_PREFIX`.
    ///
    /// Falls back to defaults when unset.
    pub fn from_embedd_env() -> Self {
        let query_prefix =
            std::env::var("EMBEDD_QUERY_PREFIX").unwrap_or_else(|_| Self::default().query_prefix);
        let doc_prefix =
            std::env::var("EMBEDD_DOC_PREFIX").unwrap_or_else(|_| Self::default().doc_prefix);
        Self {
            query_prefix,
            doc_prefix,
        }
    }

    /// Prefer `EMBEDD_*` prompt env vars, else fall back to `IKSH_*`, else defaults.
    ///
    /// This is useful for examples/benchmarks that want a single "prompt surface" without
    /// accidentally drifting existing `iksh` behavior.
    #[allow(deprecated)]
    pub fn from_env_any() -> Self {
        let has_embedd = std::env::var("EMBEDD_QUERY_PREFIX").is_ok()
            || std::env::var("EMBEDD_DOC_PREFIX").is_ok();
        if has_embedd {
            return Self::from_embedd_env();
        }

        let has_iksh = std::env::var("IKSH_EMBED_QUERY_PREFIX").is_ok()
            || std::env::var("IKSH_EMBED_DOC_PREFIX").is_ok();
        if has_iksh {
            return Self::from_iksh_env();
        }

        Self::default()
    }
}

/// Wrapper that applies a `PromptTemplate` before calling an inner `TextEmbedder`.
///
/// This is the "instruction/scoped embedding" adapter when a backend:
/// - ignores `EmbedMode`, or
/// - expects explicit prompt prefixes.
///
/// **Nuance**: If your inner backend already applies its own prompt logic (e.g. a backend that
/// uses `EmbedMode` internally), wrapping may double-prefix. Keep this opt-in.
#[derive(Debug, Clone)]
pub struct PromptedTextEmbedder<E> {
    inner: E,
    prompt: PromptTemplate,
}

impl<E> PromptedTextEmbedder<E> {
    pub fn new(inner: E, prompt: PromptTemplate) -> Self {
        Self { inner, prompt }
    }

    pub fn prompt(&self) -> &PromptTemplate {
        &self.prompt
    }

    pub fn into_inner(self) -> E {
        self.inner
    }
}

impl<E: TextEmbedder> TextEmbedder for PromptedTextEmbedder<E> {
    fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> anyhow::Result<Vec<Vec<f32>>> {
        let prompted: Vec<String> = texts.iter().map(|t| self.prompt.apply(mode, t)).collect();
        // Pass the same mode through; the inner may ignore it or use it for other behavior.
        self.inner.embed_texts(&prompted, mode)
    }

    fn model_id(&self) -> Option<&str> {
        self.inner.model_id()
    }

    fn dimension(&self) -> Option<usize> {
        self.inner.dimension()
    }

    fn capabilities(&self) -> TextEmbedderCapabilities {
        let mut caps = self.inner.capabilities();
        // This wrapper definitely applies a client-side prefix and uses EmbedMode for that.
        caps.uses_embed_mode = PromptApplication::ClientPrefix;
        caps
    }
}

/// Wrapper that enforces L2-normalized outputs.
///
/// This is the "design around normalization drift" adapter: downstream code can rely on cosine==dot.
#[derive(Debug, Clone)]
pub struct L2NormalizedTextEmbedder<E> {
    inner: E,
}

impl<E> L2NormalizedTextEmbedder<E> {
    pub fn new(inner: E) -> Self {
        Self { inner }
    }
}

impl<E: TextEmbedder> TextEmbedder for L2NormalizedTextEmbedder<E> {
    fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> anyhow::Result<Vec<Vec<f32>>> {
        let mut embs = self.inner.embed_texts(texts, mode)?;
        for e in &mut embs {
            vector::l2_normalize_in_place(e);
        }
        Ok(embs)
    }

    fn model_id(&self) -> Option<&str> {
        self.inner.model_id()
    }

    fn dimension(&self) -> Option<usize> {
        self.inner.dimension()
    }

    fn capabilities(&self) -> TextEmbedderCapabilities {
        let mut caps = self.inner.capabilities();
        caps.normalization = Normalization::L2Normalized;
        caps
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum NormalizationPolicy {
    /// Keep backend behavior.
    Preserve,
    /// Ensure outputs are L2-normalized (wrap if needed).
    RequireL2,
}

pub fn apply_normalization_policy<E: TextEmbedder + 'static>(
    inner: E,
    policy: NormalizationPolicy,
) -> anyhow::Result<Box<dyn TextEmbedder>> {
    let caps = inner.capabilities();
    match policy {
        NormalizationPolicy::Preserve => Ok(Box::new(inner)),
        NormalizationPolicy::RequireL2 => {
            if caps.normalization == Normalization::L2Normalized {
                Ok(Box::new(inner))
            } else {
                Ok(Box::new(L2NormalizedTextEmbedder::new(inner)))
            }
        }
    }
}

/// Wrapper that truncates output vectors to the first `dim` dimensions.
///
/// This matches the common "truncate_dim" / "dimensions" knob (e.g. SentenceTransformers, TEI).
#[derive(Debug, Clone)]
pub struct TruncateDimTextEmbedder<E> {
    inner: E,
    dim: usize,
}

impl<E> TruncateDimTextEmbedder<E> {
    pub fn new(inner: E, dim: usize) -> anyhow::Result<Self> {
        if dim == 0 {
            return Err(anyhow::anyhow!("truncate dim must be > 0"));
        }
        Ok(Self { inner, dim })
    }
}

impl<E: TextEmbedder> TextEmbedder for TruncateDimTextEmbedder<E> {
    fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> anyhow::Result<Vec<Vec<f32>>> {
        let mut embs = self.inner.embed_texts(texts, mode)?;
        for v in &mut embs {
            if v.len() < self.dim {
                return Err(anyhow::anyhow!(
                    "truncate dim {} exceeds embedding dim {}",
                    self.dim,
                    v.len()
                ));
            }
            v.truncate(self.dim);
        }
        Ok(embs)
    }

    fn model_id(&self) -> Option<&str> {
        self.inner.model_id()
    }

    fn dimension(&self) -> Option<usize> {
        Some(self.dim)
    }

    fn capabilities(&self) -> TextEmbedderCapabilities {
        // This wrapper doesn't claim anything about prompt/truncation tokenization behavior.
        self.inner.capabilities()
    }
}

pub fn apply_output_dim<E: TextEmbedder + 'static>(
    inner: E,
    dim: Option<usize>,
) -> anyhow::Result<Box<dyn TextEmbedder>> {
    match dim {
        None => Ok(Box::new(inner)),
        Some(d) => Ok(Box::new(TruncateDimTextEmbedder::new(inner, d)?)),
    }
}

/// Wrapper that splits large batches into fixed-size chunks before delegating.
///
/// Useful for rate-limited APIs or backends with memory constraints.
#[derive(Debug, Clone)]
pub struct BatchingTextEmbedder<E> {
    inner: E,
    batch_size: usize,
}

impl<E> BatchingTextEmbedder<E> {
    pub fn new(inner: E, batch_size: usize) -> Self {
        Self {
            inner,
            batch_size: batch_size.max(1),
        }
    }
}

impl<E: TextEmbedder> TextEmbedder for BatchingTextEmbedder<E> {
    fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> anyhow::Result<Vec<Vec<f32>>> {
        if texts.len() <= self.batch_size {
            return self.inner.embed_texts(texts, mode);
        }
        let mut out = Vec::with_capacity(texts.len());
        for chunk in texts.chunks(self.batch_size) {
            out.extend(self.inner.embed_texts(chunk, mode)?);
        }
        Ok(out)
    }

    fn model_id(&self) -> Option<&str> {
        self.inner.model_id()
    }

    fn dimension(&self) -> Option<usize> {
        self.inner.dimension()
    }

    fn capabilities(&self) -> TextEmbedderCapabilities {
        self.inner.capabilities()
    }
}

// ---------------------------------------------------------------------------
// Embedding cache wrapper
// ---------------------------------------------------------------------------

/// Wrapper that caches embedding results, avoiding redundant backend calls.
///
/// Cache key is `(model_id, mode, text)`. Texts already in the cache are returned
/// directly; only cache misses are forwarded to the inner embedder.
///
/// ```ignore
/// use embedd::{CachingTextEmbedder, EmbedMode};
///
/// let cached = CachingTextEmbedder::new(inner_embedder);
/// let v1 = cached.embed_text("hello", EmbedMode::Document)?;  // calls backend
/// let v2 = cached.embed_text("hello", EmbedMode::Document)?;  // cache hit
/// assert_eq!(v1, v2);
/// ```
pub struct CachingTextEmbedder<E> {
    inner: E,
    cache: std::sync::Mutex<std::collections::HashMap<(String, u8), Vec<f32>>>,
}

impl<E> CachingTextEmbedder<E> {
    /// Wrap an embedder with an in-memory cache.
    pub fn new(inner: E) -> Self {
        Self {
            inner,
            cache: std::sync::Mutex::new(std::collections::HashMap::new()),
        }
    }

    /// Number of cached entries.
    pub fn cache_len(&self) -> usize {
        self.cache.lock().unwrap_or_else(|e| e.into_inner()).len()
    }

    /// Clear all cached embeddings.
    pub fn clear_cache(&self) {
        self.cache.lock().unwrap_or_else(|e| e.into_inner()).clear();
    }
}

impl<E: TextEmbedder> TextEmbedder for CachingTextEmbedder<E> {
    fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> anyhow::Result<Vec<Vec<f32>>> {
        let mode_key = match mode {
            EmbedMode::Query => 0u8,
            EmbedMode::Document => 1u8,
        };

        let mut results: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
        let mut miss_indices: Vec<usize> = Vec::new();
        let mut miss_texts: Vec<String> = Vec::new();

        // Check cache for each text.
        {
            let cache = self.cache.lock().unwrap_or_else(|e| e.into_inner());
            for (i, text) in texts.iter().enumerate() {
                let key = (text.clone(), mode_key);
                if let Some(vec) = cache.get(&key) {
                    results[i] = Some(vec.clone());
                } else {
                    miss_indices.push(i);
                    miss_texts.push(text.clone());
                }
            }
        }

        // Embed cache misses.
        if !miss_texts.is_empty() {
            let embedded = self.inner.embed_texts(&miss_texts, mode)?;
            let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner());
            for (j, idx) in miss_indices.iter().enumerate() {
                if let Some(vec) = embedded.get(j) {
                    cache.insert((texts[*idx].clone(), mode_key), vec.clone());
                    results[*idx] = Some(vec.clone());
                }
            }
        }

        Ok(results.into_iter().map(|r| r.unwrap_or_default()).collect())
    }

    fn model_id(&self) -> Option<&str> {
        self.inner.model_id()
    }

    fn dimension(&self) -> Option<usize> {
        self.inner.dimension()
    }

    fn capabilities(&self) -> TextEmbedderCapabilities {
        self.inner.capabilities()
    }
}

/// Minimal interface for "text → dense vector" encoders (bi-encoder style).
///
/// This covers the common "sentence embedding" family: one vector per input string.
///
/// **Important**: there are multiple *kinds* of "embeddings" used in retrieval:
///
/// - **Dense sentence embeddings** (this trait): one \(d\)-dim vector per string.
/// - **Dense token embeddings / late interaction** (e.g. ColBERT): many vectors per string.
/// - **Sparse embeddings** (e.g. SPLADE): a weighted sparse vector over vocabulary IDs.
/// - **Binary / quantized embeddings**: compressed representations for ANN speed/memory.
///
/// We start with the dense-sentence contract because it's the smallest stable surface that
/// many parts of the workspace can share (iksh, chunking, retrieval prototypes).
pub trait TextEmbedder: Send + Sync {
    /// Embed texts into vectors.
    ///
    /// Recommended invariant for backends that can afford it: **L2-normalize** outputs so
    /// cosine similarity equals dot product and downstream fusion is less brittle.
    fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> anyhow::Result<Vec<Vec<f32>>>;

    /// Convenience: embed a single text, avoiding the `&[s.to_string()][0].clone()` boilerplate.
    fn embed_text(&self, text: &str, mode: EmbedMode) -> anyhow::Result<Vec<f32>> {
        self.embed_texts(&[text.to_string()], mode)?
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("embed_texts returned empty"))
    }

    /// Optional: model identifier for debugging and provenance.
    fn model_id(&self) -> Option<&str> {
        None
    }

    /// Optional: embedding dimension **as returned by this embedder**.
    ///
    /// Notes:
    /// - If you wrap an embedder with `apply_output_dim(Some(d))`, this should become `Some(d)`.
    /// - Many remote backends can't report this without a request; returning `None` is fine.
    fn dimension(&self) -> Option<usize> {
        None
    }

    /// Optional: backend capability declaration.
    ///
    /// This exists to design around "silent drift" failure modes:
    /// - prompt applied client-side vs server-side vs internally
    /// - whether `EmbedMode` is actually used
    /// - whether outputs are L2-normalized
    ///
    /// Backends that don't override this should be treated as "unknown" by callers.
    fn capabilities(&self) -> TextEmbedderCapabilities {
        TextEmbedderCapabilities::unknown()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum Normalization {
    Unknown,
    L2Normalized,
    NotNormalized,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum TruncationDirection {
    Unknown,
    Left,
    Right,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum TruncationPolicy {
    Unknown,
    None,
    Truncate {
        max_len: Option<usize>,
        direction: TruncationDirection,
    },
}

/// Where prompt/scoping is applied.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum PromptApplication {
    Unknown,
    None,
    /// Prefix applied in the client before sending to the model/backend.
    ClientPrefix,
    /// Prompt selection is delegated to the backend/server (e.g. TEI `prompt_name`).
    ServerPromptName,
    /// Backend applies scope/prompt internally (e.g. local HF embedder uses `EmbedMode`).
    Internal,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TextEmbedderCapabilities {
    pub uses_embed_mode: PromptApplication,
    pub normalization: Normalization,
    pub truncation: TruncationPolicy,
}

impl TextEmbedderCapabilities {
    pub const fn unknown() -> Self {
        Self {
            uses_embed_mode: PromptApplication::Unknown,
            normalization: Normalization::Unknown,
            truncation: TruncationPolicy::Unknown,
        }
    }
}

/// Declarative scoping/prompt policy for a text embedder call site.
///
/// This is intentionally small: it exists to prevent the two most common regression classes:
/// 1) **double prompting** (client prefix + server/internal prompt)
/// 2) **silent mode ignore** (caller thinks Query/Document scope matters, but backend ignores it)
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum ScopingPolicy {
    /// No scoping/prompting beyond what the backend does by default.
    None,
    /// Apply a client-side prefix template based on `EmbedMode`.
    ClientPrefix(PromptTemplate),
    /// Require that the backend applies scope via server-side prompt selection (e.g. TEI `prompt_name`).
    RequireServerPromptName,
    /// Require that the backend applies scope internally (e.g. local embedder with per-mode prompt fields).
    RequireInternal,
}

/// Apply a `ScopingPolicy` to an embedder, returning a boxed embedder.
///
/// Note: this only *wraps* for `ClientPrefix`. The other policies are validation gates.
pub fn apply_scoping_policy<E: TextEmbedder + 'static>(
    inner: E,
    policy: ScopingPolicy,
) -> anyhow::Result<Box<dyn TextEmbedder>> {
    let caps = inner.capabilities();
    match policy {
        ScopingPolicy::None => Ok(Box::new(inner)),
        ScopingPolicy::ClientPrefix(prompt) => {
            // Disallow obvious double-prompting.
            match caps.uses_embed_mode {
                PromptApplication::ServerPromptName | PromptApplication::Internal => {
                    return Err(anyhow::anyhow!(
                        "scoping policy ClientPrefix conflicts with backend prompt application {:?}",
                        caps.uses_embed_mode
                    ));
                }
                _ => {}
            }
            Ok(Box::new(PromptedTextEmbedder::new(inner, prompt)))
        }
        ScopingPolicy::RequireServerPromptName => {
            if caps.uses_embed_mode != PromptApplication::ServerPromptName {
                return Err(anyhow::anyhow!(
                    "expected ServerPromptName scoping, but backend reports {:?}",
                    caps.uses_embed_mode
                ));
            }
            Ok(Box::new(inner))
        }
        ScopingPolicy::RequireInternal => {
            if caps.uses_embed_mode != PromptApplication::Internal {
                return Err(anyhow::anyhow!(
                    "expected Internal scoping, but backend reports {:?}",
                    caps.uses_embed_mode
                ));
            }
            Ok(Box::new(inner))
        }
    }
}

// Ergonomics: allow trait objects (`Box<dyn TextEmbedder>`) to be used wherever a `TextEmbedder` is expected.
impl<T: TextEmbedder + ?Sized> TextEmbedder for Box<T> {
    fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> anyhow::Result<Vec<Vec<f32>>> {
        (**self).embed_texts(texts, mode)
    }

    fn model_id(&self) -> Option<&str> {
        (**self).model_id()
    }

    fn dimension(&self) -> Option<usize> {
        (**self).dimension()
    }

    fn capabilities(&self) -> TextEmbedderCapabilities {
        (**self).capabilities()
    }
}

/// Optional extension trait for "multi-vector" (late-interaction) embeddings.
///
/// Shape: `batch -> tokens -> dim`.
pub trait TokenEmbedder: Send + Sync {
    fn embed_tokens(&self, texts: &[String], mode: EmbedMode)
        -> anyhow::Result<Vec<Vec<Vec<f32>>>>;
}

/// Optional extension trait for sparse lexical embeddings.
///
/// Each string becomes a sparse vector: `(term_id, weight)` pairs.
pub trait SparseEmbedder: Send + Sync {
    fn embed_sparse(
        &self,
        texts: &[String],
        mode: EmbedMode,
    ) -> anyhow::Result<Vec<Vec<(u32, f32)>>>;
}

/// Minimal image embedder interface (bytes -> vectors).
///
/// `images` are opaque bytes (e.g. JPEG/PNG). Format detection/decoding is backend-specific.
pub trait ImageEmbedder: Send + Sync {
    fn embed_images(&self, images: &[Vec<u8>]) -> anyhow::Result<Vec<Vec<f32>>>;

    fn model_id(&self) -> Option<&str> {
        None
    }
}

/// Minimal audio embedder interface (bytes -> vectors).
///
/// `audios` are opaque bytes (e.g. WAV/MP3). Decoding/resampling is backend-specific.
pub trait AudioEmbedder: Send + Sync {
    fn embed_audios(&self, audios: &[Vec<u8>]) -> anyhow::Result<Vec<Vec<f32>>>;

    fn model_id(&self) -> Option<&str> {
        None
    }
}

// --- Reranker traits ---

/// A reranking result: document index paired with a relevance score.
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct RerankResult {
    /// Index of the document in the input slice.
    pub index: usize,
    /// Relevance score (higher is more relevant).
    pub score: f32,
}

/// Rerank documents by relevance to a query using a cross-encoder model.
pub trait Reranker: Send + Sync {
    /// Score and rank documents by relevance to the query.
    /// Returns results sorted by descending score.
    fn rerank(
        &self,
        query: &str,
        documents: &[String],
        top_k: Option<usize>,
    ) -> anyhow::Result<Vec<RerankResult>>;

    /// Model identifier, if available.
    fn model_id(&self) -> Option<&str> {
        None
    }
}

// Ergonomics: allow `Box<dyn Reranker>` to be used wherever a `Reranker` is expected.
impl<T: Reranker + ?Sized> Reranker for Box<T> {
    fn rerank(
        &self,
        query: &str,
        documents: &[String],
        top_k: Option<usize>,
    ) -> anyhow::Result<Vec<RerankResult>> {
        (**self).rerank(query, documents, top_k)
    }

    fn model_id(&self) -> Option<&str> {
        (**self).model_id()
    }
}

/// Wraps a `Reranker` with batch size control.
///
/// Documents are scored in chunks of `batch_size`, then merged and sorted.
/// This is useful for backends with memory constraints or rate limits.
#[derive(Debug, Clone)]
pub struct BatchingReranker<R> {
    inner: R,
    batch_size: usize,
}

impl<R> BatchingReranker<R> {
    pub fn new(inner: R, batch_size: usize) -> Self {
        Self {
            inner,
            batch_size: batch_size.max(1),
        }
    }

    pub fn into_inner(self) -> R {
        self.inner
    }
}

impl<R: Reranker> Reranker for BatchingReranker<R> {
    fn rerank(
        &self,
        query: &str,
        documents: &[String],
        top_k: Option<usize>,
    ) -> anyhow::Result<Vec<RerankResult>> {
        if documents.len() <= self.batch_size {
            return self.inner.rerank(query, documents, top_k);
        }
        let mut all_results = Vec::new();
        for (batch_idx, chunk) in documents.chunks(self.batch_size).enumerate() {
            let chunk_vec: Vec<String> = chunk.to_vec();
            let mut batch_results = self.inner.rerank(query, &chunk_vec, None)?;
            // Adjust indices to be global.
            for result in &mut batch_results {
                result.index += batch_idx * self.batch_size;
            }
            all_results.extend(batch_results);
        }
        all_results.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        if let Some(k) = top_k {
            all_results.truncate(k);
        }
        Ok(all_results)
    }

    fn model_id(&self) -> Option<&str> {
        self.inner.model_id()
    }
}

// --- Async trait ---

/// Boxed future type used by `AsyncTextEmbedder` for object safety.
#[cfg(any(
    feature = "async-openai",
    feature = "async-tei",
    feature = "async-hf-inference",
))]
pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;

/// Async counterpart of `TextEmbedder`.
///
/// Object-safe via boxed futures. Use `async-openai`, `async-tei`, or `async-hf-inference`
/// features for reqwest-backed implementations.
#[cfg(any(
    feature = "async-openai",
    feature = "async-tei",
    feature = "async-hf-inference",
))]
pub trait AsyncTextEmbedder: Send + Sync {
    fn embed_texts(
        &self,
        texts: &[String],
        mode: EmbedMode,
    ) -> BoxFuture<'_, anyhow::Result<Vec<Vec<f32>>>>;

    fn embed_text(&self, text: &str, mode: EmbedMode) -> BoxFuture<'_, anyhow::Result<Vec<f32>>> {
        let text = text.to_string();
        Box::pin(async move {
            self.embed_texts(&[text], mode)
                .await?
                .into_iter()
                .next()
                .ok_or_else(|| anyhow::anyhow!("embed_texts returned empty"))
        })
    }

    fn model_id(&self) -> Option<&str> {
        None
    }

    fn dimension(&self) -> Option<usize> {
        None
    }
}

/// Wraps any sync `TextEmbedder` as an `AsyncTextEmbedder` via `tokio::task::spawn_blocking`.
#[cfg(any(
    feature = "async-openai",
    feature = "async-tei",
    feature = "async-hf-inference",
))]
pub struct AsyncSyncBridge<E> {
    inner: std::sync::Arc<E>,
}

#[cfg(any(
    feature = "async-openai",
    feature = "async-tei",
    feature = "async-hf-inference",
))]
impl<E: TextEmbedder + 'static> AsyncSyncBridge<E> {
    pub fn new(inner: E) -> Self {
        Self {
            inner: std::sync::Arc::new(inner),
        }
    }
}

#[cfg(any(
    feature = "async-openai",
    feature = "async-tei",
    feature = "async-hf-inference",
))]
impl<E: TextEmbedder + 'static> AsyncTextEmbedder for AsyncSyncBridge<E> {
    fn embed_texts(
        &self,
        texts: &[String],
        mode: EmbedMode,
    ) -> BoxFuture<'_, anyhow::Result<Vec<Vec<f32>>>> {
        let inner = std::sync::Arc::clone(&self.inner);
        let texts = texts.to_vec();
        Box::pin(async move {
            tokio::task::spawn_blocking(move || inner.embed_texts(&texts, mode))
                .await
                .map_err(|e| anyhow::anyhow!("spawn_blocking join error: {e}"))?
        })
    }

    fn model_id(&self) -> Option<&str> {
        self.inner.model_id()
    }

    fn dimension(&self) -> Option<usize> {
        self.inner.dimension()
    }
}

/// Async counterpart of `Reranker`.
///
/// Object-safe via boxed futures. Gated on the same async features as `AsyncTextEmbedder`.
#[cfg(any(
    feature = "async-openai",
    feature = "async-tei",
    feature = "async-hf-inference",
))]
pub trait AsyncReranker: Send + Sync {
    /// Score and rank documents by relevance to the query.
    fn rerank(
        &self,
        query: &str,
        documents: &[String],
        top_k: Option<usize>,
    ) -> BoxFuture<'_, anyhow::Result<Vec<RerankResult>>>;

    /// Model identifier, if available.
    fn model_id(&self) -> Option<&str> {
        None
    }
}

// --- Optional backends (feature-gated) ---
//
// Goal: `embedd` is the only public crate name. Backends live behind features,
// not as separate publishable crates.

#[cfg(feature = "hf-inference")]
pub mod hf_inference {
    use super::{AudioEmbedder, ImageEmbedder, TextEmbedder};
    use anyhow::Result;

    /// HuggingFace Inference API "feature-extraction" embedder.
    ///
    /// This is a **real** multimodal e2e path (network), useful for:
    /// - quickly smoke-testing text/image/audio modality plumbing
    /// - comparing remote models without pulling weights locally
    ///
    /// Auth: set `HF_API_TOKEN` env var, or pass the token explicitly.
    ///
    /// Endpoint shape is assumed to be `POST https://api-inference.huggingface.co/models/{model}`.
    /// The response is expected to be numeric arrays (often nested). We reduce to a single vector
    /// via mean pooling over leading dimensions.
    #[derive(Debug, Clone)]
    pub struct HfInferenceEmbedder {
        base_url: String,
        token: Option<String>,
        model: String,
        client: ureq::Agent,
    }

    impl HfInferenceEmbedder {
        pub fn new(model: impl Into<String>) -> Self {
            Self::new_with_base_url(model, "https://api-inference.huggingface.co")
        }

        pub fn new_with_base_url(model: impl Into<String>, base_url: impl Into<String>) -> Self {
            Self {
                base_url: base_url.into().trim_end_matches('/').to_string(),
                token: std::env::var("HF_API_TOKEN").ok(),
                model: model.into(),
                client: ureq::AgentBuilder::new().build(),
            }
        }

        pub fn with_token(mut self, token: impl Into<String>) -> Self {
            self.token = Some(token.into());
            self
        }

        pub fn model(&self) -> &str {
            &self.model
        }

        fn endpoint(&self) -> String {
            format!("{}/models/{}", self.base_url, self.model)
        }

        fn auth_header_value(&self) -> Option<String> {
            self.token.as_ref().map(|t| format!("Bearer {t}"))
        }

        fn post_json<T: serde::Serialize>(&self, payload: &T) -> Result<serde_json::Value> {
            let mut req = self.client.post(&self.endpoint());
            if let Some(auth) = self.auth_header_value() {
                req = req.set("Authorization", &auth);
            }

            // ureq 2.x: send_json handles Content-Type + serialization; Err on non-2xx.
            let resp = req.send_json(payload)?;
            let body = resp.into_string().unwrap_or_default();
            Ok(serde_json::from_str(&body)?)
        }

        fn post_bytes(&self, bytes: &[u8]) -> Result<serde_json::Value> {
            let mut req = self
                .client
                .post(&self.endpoint())
                .set("Content-Type", "application/octet-stream");
            if let Some(auth) = self.auth_header_value() {
                req = req.set("Authorization", &auth);
            }

            // ureq 2.x returns Err for non-2xx.
            let resp = req.send_bytes(bytes)?;
            let body = resp.into_string().unwrap_or_default();
            Ok(serde_json::from_str(&body)?)
        }
    }

    fn mean_pool_to_vec(v: &serde_json::Value) -> Result<Vec<f32>> {
        // Accept shapes like:
        // - [d]
        // - [[d]]
        // - [[[d]]] (token/patch embeddings)
        // and mean-pool all but the last dimension.
        fn flatten_numbers(v: &serde_json::Value, out: &mut Vec<f64>) -> Result<()> {
            match v {
                serde_json::Value::Number(n) => {
                    out.push(
                        n.as_f64()
                            .ok_or_else(|| anyhow::anyhow!("non-f64 number"))?,
                    );
                    Ok(())
                }
                serde_json::Value::Array(xs) => {
                    for x in xs {
                        flatten_numbers(x, out)?;
                    }
                    Ok(())
                }
                _ => Err(anyhow::anyhow!("expected numeric arrays, got: {}", v)),
            }
        }

        // Strategy:
        // - If it's a flat vector already, return it.
        // - Otherwise: interpret the last dimension as `d`, average across leading elements.
        fn as_vec_f32(v: &serde_json::Value) -> Option<Vec<f32>> {
            match v {
                serde_json::Value::Array(xs) if xs.iter().all(|x| x.is_number()) => xs
                    .iter()
                    .map(|x| x.as_f64().map(|f| f as f32))
                    .collect::<Option<Vec<f32>>>(),
                _ => None,
            }
        }

        if let Some(vec) = as_vec_f32(v) {
            return Ok(vec);
        }

        // Try to treat v as nested arrays whose leaves are vectors of the same dimension.
        fn collect_leaf_vectors(v: &serde_json::Value, leaves: &mut Vec<Vec<f32>>) -> Result<()> {
            if let Some(vec) = as_vec_f32(v) {
                leaves.push(vec);
                return Ok(());
            }
            match v {
                serde_json::Value::Array(xs) => {
                    for x in xs {
                        collect_leaf_vectors(x, leaves)?;
                    }
                    Ok(())
                }
                _ => Err(anyhow::anyhow!("expected numeric arrays, got: {}", v)),
            }
        }

        let mut leaves = Vec::new();
        collect_leaf_vectors(v, &mut leaves)?;
        if leaves.is_empty() {
            // fallback: full flatten (shouldn't happen for feature-extraction, but keep a clear error)
            let mut flat = Vec::new();
            flatten_numbers(v, &mut flat)?;
            return Ok(flat.into_iter().map(|x| x as f32).collect());
        }

        let d = leaves[0].len();
        if d == 0 {
            return Err(anyhow::anyhow!("zero-dim embedding"));
        }
        if !leaves.iter().all(|e| e.len() == d) {
            return Err(anyhow::anyhow!(
                "inconsistent embedding dims in HF response"
            ));
        }

        let mut acc = vec![0.0f32; d];
        for e in &leaves {
            for (i, x) in e.iter().enumerate() {
                acc[i] += *x;
            }
        }
        let n = leaves.len() as f32;
        for x in &mut acc {
            *x /= n;
        }
        Ok(acc)
    }

    impl TextEmbedder for HfInferenceEmbedder {
        fn embed_texts(&self, texts: &[String], _mode: super::EmbedMode) -> Result<Vec<Vec<f32>>> {
            let payload = serde_json::json!({ "inputs": texts });
            let v = self.post_json(&payload)?;
            // Response is typically: Vec<Vec<f32>> (batch -> vector), but can vary.
            // We accept:
            // - [ [d], [d], ... ]
            // - [d] (single input) -> wrap
            match v {
                serde_json::Value::Array(items) => {
                    // If items are numbers, treat as single embedding vector.
                    if items.iter().all(|x| x.is_number()) {
                        return Ok(vec![mean_pool_to_vec(&serde_json::Value::Array(items))?]);
                    }
                    // Otherwise, each item should be a (possibly nested) embedding.
                    let mut out = Vec::with_capacity(items.len());
                    for item in items {
                        out.push(mean_pool_to_vec(&item)?);
                    }
                    Ok(out)
                }
                other => Err(anyhow::anyhow!("unexpected HF response: {}", other)),
            }
        }

        fn model_id(&self) -> Option<&str> {
            Some(&self.model)
        }

        fn capabilities(&self) -> super::TextEmbedderCapabilities {
            super::TextEmbedderCapabilities {
                uses_embed_mode: super::PromptApplication::None,
                normalization: super::Normalization::Unknown,
                truncation: super::TruncationPolicy::Unknown,
            }
        }
    }

    impl ImageEmbedder for HfInferenceEmbedder {
        fn embed_images(&self, images: &[Vec<u8>]) -> Result<Vec<Vec<f32>>> {
            let mut out = Vec::with_capacity(images.len());
            for img in images {
                let v = self.post_bytes(img)?;
                out.push(mean_pool_to_vec(&v)?);
            }
            Ok(out)
        }

        fn model_id(&self) -> Option<&str> {
            Some(&self.model)
        }
    }

    impl AudioEmbedder for HfInferenceEmbedder {
        fn embed_audios(&self, audios: &[Vec<u8>]) -> Result<Vec<Vec<f32>>> {
            let mut out = Vec::with_capacity(audios.len());
            for a in audios {
                let v = self.post_bytes(a)?;
                out.push(mean_pool_to_vec(&v)?);
            }
            Ok(out)
        }

        fn model_id(&self) -> Option<&str> {
            Some(&self.model)
        }
    }
}

#[cfg(feature = "openai")]
pub mod openai {
    use super::{EmbedMode, TextEmbedder};
    use anyhow::Result;
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Clone)]
    pub struct OpenAiEmbedder {
        base_url: String,
        api_key: String,
        model: String,
        client: ureq::Agent,
    }

    impl OpenAiEmbedder {
        /// Create an embedder targeting the OpenAI API (default base URL).
        pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
            Self {
                base_url: "https://api.openai.com".to_string(),
                api_key: api_key.into(),
                model: model.into(),
                client: ureq::AgentBuilder::new().build(),
            }
        }

        /// Override the base URL (for OpenAI-compatible APIs like vLLM, Ollama, etc.).
        pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
            self.base_url = base_url.into().trim_end_matches('/').to_string();
            self
        }

        fn embeddings_endpoint(&self) -> String {
            format!("{}/v1/embeddings", self.base_url)
        }
    }

    #[derive(Debug, Serialize)]
    struct EmbeddingsRequest<'a> {
        model: &'a str,
        input: &'a [String],
    }

    #[derive(Debug, Deserialize)]
    struct EmbeddingsResponse {
        data: Vec<EmbeddingDatum>,
    }

    #[derive(Debug, Deserialize)]
    struct EmbeddingDatum {
        embedding: Vec<f32>,
    }

    impl TextEmbedder for OpenAiEmbedder {
        fn embed_texts(&self, texts: &[String], _mode: EmbedMode) -> Result<Vec<Vec<f32>>> {
            let payload = EmbeddingsRequest {
                model: &self.model,
                input: texts,
            };

            // ureq 2.x returns Err for non-2xx, so we only reach here on success.
            let resp = self
                .client
                .post(&self.embeddings_endpoint())
                .set("Authorization", &format!("Bearer {}", self.api_key))
                .send_json(&payload)?;

            let body = resp.into_string()?;
            let parsed: EmbeddingsResponse = serde_json::from_str(&body)?;
            Ok(parsed.data.into_iter().map(|d| d.embedding).collect())
        }

        fn model_id(&self) -> Option<&str> {
            Some(&self.model)
        }

        fn capabilities(&self) -> super::TextEmbedderCapabilities {
            super::TextEmbedderCapabilities {
                uses_embed_mode: super::PromptApplication::None,
                normalization: super::Normalization::Unknown,
                truncation: super::TruncationPolicy::Unknown,
            }
        }
    }
}

#[cfg(feature = "tei")]
pub mod tei {
    use super::{
        EmbedMode, PromptApplication, TextEmbedder, TextEmbedderCapabilities, TruncationDirection,
        TruncationPolicy,
    };
    use anyhow::Result;
    use serde_json::Value;

    /// TEI client embedder.
    #[derive(Debug, Clone)]
    pub struct TeiEmbedder {
        base_url: String,
        api_key: Option<String>,
        prompt_name_query: Option<String>,
        prompt_name_doc: Option<String>,
        dimensions: Option<usize>,
        normalize: Option<bool>,
        truncate: Option<bool>,
        truncation_direction: Option<String>,
        client: ureq::Agent,
    }

    impl TeiEmbedder {
        /// Create a new TEI client targeting `base_url` (e.g. `http://127.0.0.1:8080`).
        pub fn new(base_url: impl Into<String>) -> Self {
            Self {
                base_url: base_url.into().trim_end_matches('/').to_string(),
                api_key: None,
                prompt_name_query: None,
                prompt_name_doc: None,
                dimensions: None,
                normalize: None,
                truncate: None,
                truncation_direction: None,
                client: ureq::AgentBuilder::new().build(),
            }
        }

        /// Set a TEI API key (TEI supports an `--api-key` flag; this sends `Authorization: Bearer <key>`).
        pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
            self.api_key = Some(api_key.into());
            self
        }

        /// Configure TEI `prompt_name` per scope (`EmbedMode`).
        ///
        /// TEI applies prompts server-side based on the model's SentenceTransformers `prompts` config.
        pub fn with_prompt_names(
            mut self,
            query_prompt_name: impl Into<String>,
            doc_prompt_name: impl Into<String>,
        ) -> Self {
            self.prompt_name_query = Some(query_prompt_name.into());
            self.prompt_name_doc = Some(doc_prompt_name.into());
            self
        }

        /// TEI request option: force output embedding dimension (TEI `dimensions`).
        pub fn with_dimensions(mut self, dimensions: usize) -> Self {
            self.dimensions = Some(dimensions);
            self
        }

        /// TEI request option: whether to L2-normalize outputs (TEI defaults `true`).
        pub fn with_normalize(mut self, normalize: bool) -> Self {
            self.normalize = Some(normalize);
            self
        }

        /// TEI request option: whether to truncate inputs instead of erroring (TEI defaults `false`).
        pub fn with_truncate(mut self, truncate: bool) -> Self {
            self.truncate = Some(truncate);
            self
        }

        /// TEI request option: truncation direction ("left" or "right"; TEI defaults "right").
        pub fn with_truncation_direction(mut self, dir: impl Into<String>) -> Self {
            self.truncation_direction = Some(dir.into());
            self
        }

        fn prompt_name_for_mode(&self, mode: EmbedMode) -> Option<&str> {
            match mode {
                EmbedMode::Query => self.prompt_name_query.as_deref(),
                EmbedMode::Document => self.prompt_name_doc.as_deref(),
            }
        }

        fn embed_endpoint(&self) -> String {
            format!("{}/embed", self.base_url)
        }

        fn build_embed_payload(&self, texts: &[String], mode: EmbedMode) -> Value {
            let mut payload = serde_json::Map::new();
            payload.insert("inputs".to_string(), serde_json::Value::from(texts));
            if let Some(d) = self.dimensions {
                payload.insert("dimensions".to_string(), serde_json::Value::from(d as u64));
            }
            if let Some(pn) = self.prompt_name_for_mode(mode) {
                payload.insert("prompt_name".to_string(), serde_json::Value::from(pn));
            }
            if let Some(n) = self.normalize {
                payload.insert("normalize".to_string(), serde_json::Value::from(n));
            }
            if let Some(t) = self.truncate {
                payload.insert("truncate".to_string(), serde_json::Value::from(t));
            }
            if let Some(dir) = self.truncation_direction.as_deref() {
                payload.insert(
                    "truncation_direction".to_string(),
                    serde_json::Value::from(dir),
                );
            }
            Value::Object(payload)
        }
    }

    impl TextEmbedder for TeiEmbedder {
        fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> Result<Vec<Vec<f32>>> {
            let payload = self.build_embed_payload(texts, mode);

            let mut req = self.client.post(&self.embed_endpoint());
            if let Some(k) = &self.api_key {
                req = req.set("Authorization", &format!("Bearer {k}"));
            }
            // ureq 2.x: send_json handles Content-Type + serialization; Err on non-2xx.
            let resp = req.send_json(&payload)?;
            let body = resp.into_string()?;
            let embs: Vec<Vec<f32>> = serde_json::from_str(&body)?;
            Ok(embs)
        }

        fn dimension(&self) -> Option<usize> {
            self.dimensions
        }

        fn capabilities(&self) -> TextEmbedderCapabilities {
            // TEI can apply prompts server-side if `prompt_name_*` is configured; otherwise it ignores mode.
            let uses = if self.prompt_name_query.is_some() || self.prompt_name_doc.is_some() {
                PromptApplication::ServerPromptName
            } else {
                PromptApplication::None
            };
            let truncation = match self.truncate {
                Some(true) => TruncationPolicy::Truncate {
                    max_len: None,
                    direction: match self.truncation_direction.as_deref() {
                        Some("left") => TruncationDirection::Left,
                        Some("right") => TruncationDirection::Right,
                        Some(_) => TruncationDirection::Unknown,
                        None => TruncationDirection::Right,
                    },
                },
                _ => TruncationPolicy::None,
            };
            let normalization = match self.normalize {
                Some(true) => super::Normalization::L2Normalized,
                Some(false) => super::Normalization::NotNormalized,
                None => super::Normalization::Unknown,
            };
            TextEmbedderCapabilities {
                uses_embed_mode: uses,
                normalization,
                truncation,
            }
        }
    }

    #[cfg(test)]
    mod tests {
        use super::TeiEmbedder;
        use crate::{
            EmbedMode, Normalization, PromptApplication, TextEmbedder, TruncationDirection,
            TruncationPolicy,
        };

        fn obj(v: serde_json::Value) -> serde_json::Map<String, serde_json::Value> {
            v.as_object().cloned().expect("expected object")
        }

        #[test]
        fn tei_payload_includes_prompt_name_per_mode() {
            let e = TeiEmbedder::new("http://127.0.0.1:8080").with_prompt_names("q", "d");
            let xs = vec!["hello".to_string()];

            let q = obj(e.build_embed_payload(&xs, EmbedMode::Query));
            assert_eq!(q.get("prompt_name").and_then(|v| v.as_str()), Some("q"));

            let d = obj(e.build_embed_payload(&xs, EmbedMode::Document));
            assert_eq!(d.get("prompt_name").and_then(|v| v.as_str()), Some("d"));
        }

        #[test]
        fn tei_payload_includes_dimensions_normalize_truncate_fields_when_set() {
            let e = TeiEmbedder::new("http://127.0.0.1:8080")
                .with_dimensions(256)
                .with_normalize(false)
                .with_truncate(true)
                .with_truncation_direction("left");
            let xs = vec!["hello".to_string(), "world".to_string()];
            let p = obj(e.build_embed_payload(&xs, EmbedMode::Query));

            assert_eq!(p.get("dimensions").and_then(|v| v.as_u64()), Some(256));
            assert_eq!(p.get("normalize").and_then(|v| v.as_bool()), Some(false));
            assert_eq!(p.get("truncate").and_then(|v| v.as_bool()), Some(true));
            assert_eq!(
                p.get("truncation_direction").and_then(|v| v.as_str()),
                Some("left")
            );
        }

        #[test]
        fn tei_capabilities_reflect_prompt_and_truncation_and_normalization() {
            let e = TeiEmbedder::new("http://127.0.0.1:8080")
                .with_prompt_names("query", "doc")
                .with_normalize(true)
                .with_truncate(true)
                .with_truncation_direction("right");
            let caps = e.capabilities();

            assert_eq!(caps.uses_embed_mode, PromptApplication::ServerPromptName);
            assert_eq!(caps.normalization, Normalization::L2Normalized);
            assert_eq!(
                caps.truncation,
                TruncationPolicy::Truncate {
                    max_len: None,
                    direction: TruncationDirection::Right
                }
            );
        }
    }
}

#[cfg(feature = "fastembed")]
#[path = "fastembed.rs"]
pub mod fastembed;

#[cfg(feature = "ort-tokenizers")]
#[path = "ort.rs"]
pub mod ort;

// --- Async backends ---

#[cfg(feature = "async-openai")]
pub mod async_openai {
    use super::{AsyncTextEmbedder, EmbedMode};
    use anyhow::Result;
    use serde::{Deserialize, Serialize};

    #[derive(Debug, Clone)]
    pub struct AsyncOpenAiEmbedder {
        base_url: String,
        api_key: String,
        model: String,
        client: reqwest::Client,
    }

    impl AsyncOpenAiEmbedder {
        pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
            Self {
                base_url: "https://api.openai.com".to_string(),
                api_key: api_key.into(),
                model: model.into(),
                client: reqwest::Client::new(),
            }
        }

        pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
            self.base_url = base_url.into().trim_end_matches('/').to_string();
            self
        }

        pub fn with_client(mut self, client: reqwest::Client) -> Self {
            self.client = client;
            self
        }

        fn endpoint(&self) -> String {
            format!("{}/v1/embeddings", self.base_url)
        }
    }

    #[derive(Serialize)]
    struct EmbeddingsRequest<'a> {
        model: &'a str,
        input: &'a [String],
    }

    #[derive(Deserialize)]
    struct EmbeddingsResponse {
        data: Vec<EmbeddingDatum>,
    }

    #[derive(Deserialize)]
    struct EmbeddingDatum {
        embedding: Vec<f32>,
    }

    impl AsyncTextEmbedder for AsyncOpenAiEmbedder {
        fn embed_texts(
            &self,
            texts: &[String],
            _mode: EmbedMode,
        ) -> super::BoxFuture<'_, Result<Vec<Vec<f32>>>> {
            let texts = texts.to_vec();
            Box::pin(async move {
                let payload = EmbeddingsRequest {
                    model: &self.model,
                    input: &texts,
                };
                let resp = self
                    .client
                    .post(self.endpoint())
                    .bearer_auth(&self.api_key)
                    .json(&payload)
                    .send()
                    .await?;
                let status = resp.status();
                if !status.is_success() {
                    let body = resp.text().await.unwrap_or_default();
                    return Err(anyhow::anyhow!(
                        "openai embeddings failed: status={} body={}",
                        status,
                        body
                    ));
                }
                let parsed: EmbeddingsResponse = resp.json().await?;
                Ok(parsed.data.into_iter().map(|d| d.embedding).collect())
            })
        }

        fn model_id(&self) -> Option<&str> {
            Some(&self.model)
        }
    }
}

#[cfg(feature = "async-tei")]
pub mod async_tei {
    use super::{AsyncTextEmbedder, EmbedMode};
    use anyhow::Result;

    #[derive(Debug, Clone)]
    pub struct AsyncTeiEmbedder {
        base_url: String,
        api_key: Option<String>,
        prompt_name_query: Option<String>,
        prompt_name_doc: Option<String>,
        dimensions: Option<usize>,
        normalize: Option<bool>,
        truncate: Option<bool>,
        truncation_direction: Option<String>,
        client: reqwest::Client,
    }

    impl AsyncTeiEmbedder {
        pub fn new(base_url: impl Into<String>) -> Self {
            Self {
                base_url: base_url.into().trim_end_matches('/').to_string(),
                api_key: None,
                prompt_name_query: None,
                prompt_name_doc: None,
                dimensions: None,
                normalize: None,
                truncate: None,
                truncation_direction: None,
                client: reqwest::Client::new(),
            }
        }

        pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
            self.api_key = Some(api_key.into());
            self
        }

        pub fn with_prompt_names(
            mut self,
            query: impl Into<String>,
            doc: impl Into<String>,
        ) -> Self {
            self.prompt_name_query = Some(query.into());
            self.prompt_name_doc = Some(doc.into());
            self
        }

        pub fn with_dimensions(mut self, d: usize) -> Self {
            self.dimensions = Some(d);
            self
        }

        pub fn with_normalize(mut self, n: bool) -> Self {
            self.normalize = Some(n);
            self
        }

        pub fn with_truncate(mut self, t: bool) -> Self {
            self.truncate = Some(t);
            self
        }

        pub fn with_truncation_direction(mut self, dir: impl Into<String>) -> Self {
            self.truncation_direction = Some(dir.into());
            self
        }

        pub fn with_client(mut self, client: reqwest::Client) -> Self {
            self.client = client;
            self
        }

        fn embed_endpoint(&self) -> String {
            format!("{}/embed", self.base_url)
        }

        fn build_payload(&self, texts: &[String], mode: EmbedMode) -> serde_json::Value {
            let mut m = serde_json::Map::new();
            m.insert("inputs".to_string(), serde_json::Value::from(texts));
            if let Some(d) = self.dimensions {
                m.insert("dimensions".to_string(), (d as u64).into());
            }
            let pn = match mode {
                EmbedMode::Query => self.prompt_name_query.as_deref(),
                EmbedMode::Document => self.prompt_name_doc.as_deref(),
            };
            if let Some(pn) = pn {
                m.insert("prompt_name".to_string(), pn.into());
            }
            if let Some(n) = self.normalize {
                m.insert("normalize".to_string(), n.into());
            }
            if let Some(t) = self.truncate {
                m.insert("truncate".to_string(), t.into());
            }
            if let Some(dir) = &self.truncation_direction {
                m.insert("truncation_direction".to_string(), dir.clone().into());
            }
            serde_json::Value::Object(m)
        }
    }

    impl AsyncTextEmbedder for AsyncTeiEmbedder {
        fn embed_texts(
            &self,
            texts: &[String],
            mode: EmbedMode,
        ) -> super::BoxFuture<'_, Result<Vec<Vec<f32>>>> {
            let payload = self.build_payload(texts, mode);
            Box::pin(async move {
                let mut req = self.client.post(self.embed_endpoint());
                if let Some(k) = &self.api_key {
                    req = req.bearer_auth(k);
                }
                let resp = req.json(&payload).send().await?;
                let status = resp.status();
                if !status.is_success() {
                    let body = resp.text().await.unwrap_or_default();
                    return Err(anyhow::anyhow!(
                        "tei /embed failed: status={} body={}",
                        status,
                        body
                    ));
                }
                let embs: Vec<Vec<f32>> = resp.json().await?;
                Ok(embs)
            })
        }

        fn dimension(&self) -> Option<usize> {
            self.dimensions
        }
    }
}

#[cfg(feature = "async-hf-inference")]
pub mod async_hf_inference {
    use super::{AsyncTextEmbedder, EmbedMode};
    use anyhow::Result;

    #[derive(Debug, Clone)]
    pub struct AsyncHfInferenceEmbedder {
        base_url: String,
        token: Option<String>,
        model: String,
        client: reqwest::Client,
    }

    impl AsyncHfInferenceEmbedder {
        pub fn new(model: impl Into<String>) -> Self {
            Self::new_with_base_url(model, "https://api-inference.huggingface.co")
        }

        pub fn new_with_base_url(model: impl Into<String>, base_url: impl Into<String>) -> Self {
            Self {
                base_url: base_url.into().trim_end_matches('/').to_string(),
                token: std::env::var("HF_API_TOKEN").ok(),
                model: model.into(),
                client: reqwest::Client::new(),
            }
        }

        pub fn with_token(mut self, token: impl Into<String>) -> Self {
            self.token = Some(token.into());
            self
        }

        pub fn with_client(mut self, client: reqwest::Client) -> Self {
            self.client = client;
            self
        }

        fn endpoint(&self) -> String {
            format!("{}/models/{}", self.base_url, self.model)
        }
    }

    /// Flatten nested JSON arrays to a single f32 vector via mean pooling.
    fn mean_pool(v: &serde_json::Value) -> Result<Vec<f32>> {
        fn as_flat(v: &serde_json::Value) -> Option<Vec<f32>> {
            match v {
                serde_json::Value::Array(xs) if xs.iter().all(|x| x.is_number()) => {
                    xs.iter().map(|x| x.as_f64().map(|f| f as f32)).collect()
                }
                _ => None,
            }
        }

        if let Some(vec) = as_flat(v) {
            return Ok(vec);
        }

        fn collect_leaves(v: &serde_json::Value, out: &mut Vec<Vec<f32>>) -> Result<()> {
            if let Some(vec) = as_flat(v) {
                out.push(vec);
                return Ok(());
            }
            match v {
                serde_json::Value::Array(xs) => {
                    for x in xs {
                        collect_leaves(x, out)?;
                    }
                    Ok(())
                }
                _ => Err(anyhow::anyhow!("expected numeric arrays")),
            }
        }

        let mut leaves = Vec::new();
        collect_leaves(v, &mut leaves)?;
        if leaves.is_empty() {
            return Err(anyhow::anyhow!("empty embedding response"));
        }
        let d = leaves[0].len();
        if d == 0 || !leaves.iter().all(|l| l.len() == d) {
            return Err(anyhow::anyhow!("inconsistent dims in HF response"));
        }
        let mut acc = vec![0.0f32; d];
        for leaf in &leaves {
            for (i, x) in leaf.iter().enumerate() {
                acc[i] += x;
            }
        }
        let n = leaves.len() as f32;
        for x in &mut acc {
            *x /= n;
        }
        Ok(acc)
    }

    impl AsyncTextEmbedder for AsyncHfInferenceEmbedder {
        fn embed_texts(
            &self,
            texts: &[String],
            _mode: EmbedMode,
        ) -> super::BoxFuture<'_, Result<Vec<Vec<f32>>>> {
            let payload = serde_json::json!({ "inputs": texts.to_vec() });
            Box::pin(async move {
                let mut req = self.client.post(self.endpoint()).json(&payload);
                if let Some(t) = &self.token {
                    req = req.bearer_auth(t);
                }
                let resp = req.send().await?;
                let status = resp.status();
                if !status.is_success() {
                    let body = resp.text().await.unwrap_or_default();
                    return Err(anyhow::anyhow!(
                        "hf-inference failed: status={} body={}",
                        status,
                        body
                    ));
                }
                let v: serde_json::Value = resp.json().await?;
                match v {
                    serde_json::Value::Array(items) => {
                        if items.iter().all(|x| x.is_number()) {
                            return Ok(vec![mean_pool(&serde_json::Value::Array(items))?]);
                        }
                        items.iter().map(mean_pool).collect()
                    }
                    other => Err(anyhow::anyhow!("unexpected HF response: {}", other)),
                }
            })
        }

        fn model_id(&self) -> Option<&str> {
            Some(&self.model)
        }
    }
}

/// Configuration of where embedding model artifacts come from.
///
/// This is intentionally minimal and testable: it just answers "local dir or hub id?".
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum ModelSource {
    /// Load from a local directory containing `config.json`, `tokenizer.json`, `model.safetensors`.
    LocalDir(std::path::PathBuf),
    /// Load via HuggingFace Hub model id (downloaded/cached by `hf-hub`).
    HuggingFaceModelId(String),
}

impl ModelSource {
    const DEFAULT_MODEL: &'static str = "BAAI/bge-small-en-v1.5";

    /// Resolve from the legacy `iksh` environment variables.
    ///
    /// Priority:
    /// 1) `IKSH_EMBED_MODEL_DIR` (local, no network)
    /// 2) `IKSH_EMBED_MODEL` (hub model id)
    #[deprecated(since = "0.2.0", note = "use from_embedd_env() or from_env_any()")]
    pub fn from_iksh_env() -> Self {
        if let Ok(dir) = std::env::var("IKSH_EMBED_MODEL_DIR") {
            return Self::LocalDir(std::path::PathBuf::from(dir));
        }
        let model_id =
            std::env::var("IKSH_EMBED_MODEL").unwrap_or_else(|_| Self::DEFAULT_MODEL.to_string());
        Self::HuggingFaceModelId(model_id)
    }

    /// Resolve from `EMBEDD_MODEL_DIR` / `EMBEDD_MODEL` env vars.
    ///
    /// Falls back to the default model when unset.
    pub fn from_embedd_env() -> Self {
        if let Ok(dir) = std::env::var("EMBEDD_MODEL_DIR") {
            return Self::LocalDir(std::path::PathBuf::from(dir));
        }
        let model_id =
            std::env::var("EMBEDD_MODEL").unwrap_or_else(|_| Self::DEFAULT_MODEL.to_string());
        Self::HuggingFaceModelId(model_id)
    }

    /// Prefer `EMBEDD_*` model env vars, else fall back to `IKSH_*`, else defaults.
    #[allow(deprecated)]
    pub fn from_env_any() -> Self {
        let has_embedd =
            std::env::var("EMBEDD_MODEL_DIR").is_ok() || std::env::var("EMBEDD_MODEL").is_ok();
        if has_embedd {
            return Self::from_embedd_env();
        }

        let has_iksh = std::env::var("IKSH_EMBED_MODEL_DIR").is_ok()
            || std::env::var("IKSH_EMBED_MODEL").is_ok();
        if has_iksh {
            return Self::from_iksh_env();
        }

        Self::HuggingFaceModelId(Self::DEFAULT_MODEL.to_string())
    }
}

/// Safetensors validation utilities.
///
/// This is a "native safety rail" inspired by:
/// - HuggingFace `safetensors` (Rust) invariants (bounds, contiguity, overflow checks)
/// - tinygrad's minimal `safe_load` implementation (header length + JSON + offsets)
pub mod safetensors {
    use anyhow::Result;
    use serde_json::Value;
    use std::fs::File;
    use std::io::Read;
    use std::path::Path;

    /// Hard cap to avoid allocating absurd headers.
    ///
    /// HF safetensors uses 100MB; we keep the same default.
    pub const MAX_HEADER_SIZE: usize = 100_000_000;

    #[derive(Debug, Clone, PartialEq, Eq)]
    pub struct TensorInfo {
        pub dtype: String,
        pub shape: Vec<usize>,
        /// (begin, end) offsets in bytes, relative to the data buffer start.
        pub data_offsets: (usize, usize),
    }

    fn dtype_size_bytes(dtype: &str) -> Option<usize> {
        // Conservative: cover common safetensors dtypes. Unknown dtypes fail validation.
        Some(match dtype {
            "BOOL" => 1,
            "U8" => 1,
            "I8" => 1,
            "U16" => 2,
            "I16" => 2,
            "F16" => 2,
            "BF16" => 2,
            "U32" => 4,
            "I32" => 4,
            "F32" => 4,
            "U64" => 8,
            "I64" => 8,
            "F64" => 8,
            _ => return None,
        })
    }

    fn checked_numel(shape: &[usize]) -> Option<usize> {
        let mut n: usize = 1;
        for &d in shape {
            n = n.checked_mul(d)?;
        }
        Some(n)
    }

    fn parse_usize(v: &Value) -> Option<usize> {
        v.as_u64().and_then(|x| usize::try_from(x).ok())
    }

    fn parse_tensor_info(v: &Value) -> Result<TensorInfo> {
        let obj = v
            .as_object()
            .ok_or_else(|| anyhow::anyhow!("tensor entry must be an object"))?;

        let dtype = obj
            .get("dtype")
            .and_then(|x| x.as_str())
            .ok_or_else(|| anyhow::anyhow!("tensor entry missing dtype"))?
            .to_string();

        let shape = obj
            .get("shape")
            .and_then(|x| x.as_array())
            .ok_or_else(|| anyhow::anyhow!("tensor entry missing shape"))?
            .iter()
            .map(|x| parse_usize(x).ok_or_else(|| anyhow::anyhow!("shape must be u64 array")))
            .collect::<Result<Vec<usize>>>()?;

        let offs = obj
            .get("data_offsets")
            .and_then(|x| x.as_array())
            .ok_or_else(|| anyhow::anyhow!("tensor entry missing data_offsets"))?;
        if offs.len() != 2 {
            return Err(anyhow::anyhow!("data_offsets must have length 2"));
        }
        let begin = parse_usize(&offs[0]).ok_or_else(|| anyhow::anyhow!("offset begin invalid"))?;
        let end = parse_usize(&offs[1]).ok_or_else(|| anyhow::anyhow!("offset end invalid"))?;

        Ok(TensorInfo {
            dtype,
            shape,
            data_offsets: (begin, end),
        })
    }

    /// Validate a safetensors header JSON blob against a given data buffer length.
    ///
    /// This checks:
    /// - header is valid JSON object
    /// - each tensor has known dtype, shape, and offsets
    /// - per-tensor byte size matches dtype*shape
    /// - offsets are within bounds, non-overlapping, and contiguous (no holes)
    /// - final end offset equals `data_len` (full coverage, no trailing junk)
    pub fn validate_header_and_data_len(header_bytes: &[u8], data_len: usize) -> Result<()> {
        if header_bytes.is_empty() {
            return Err(anyhow::anyhow!("header is empty"));
        }
        if header_bytes[0] != b'{' {
            return Err(anyhow::anyhow!("header must start with '{{'"));
        }

        // safetensors allows trailing whitespace padding.
        let mut trimmed = header_bytes;
        while trimmed.last() == Some(&b' ') {
            trimmed = &trimmed[..trimmed.len() - 1];
        }

        let root: Value = serde_json::from_slice(trimmed)
            .map_err(|e| anyhow::anyhow!("invalid header JSON: {e}"))?;
        let obj = root
            .as_object()
            .ok_or_else(|| anyhow::anyhow!("header must be a JSON object"))?;

        // Collect tensor infos and sort by begin offset.
        //
        // Note: empty tensors (0 bytes) are allowed by safetensors. That means multiple tensors
        // may legitimately share the same `(begin=end)` offset. We therefore cannot key by
        // `begin` alone.
        let mut tensors: Vec<(usize, usize, String, TensorInfo)> = Vec::new();
        for (k, v) in obj {
            if k == "__metadata__" {
                // safetensors restricts metadata to string->string, but we treat it as opaque.
                continue;
            }
            let info = parse_tensor_info(v)?;

            // dtype must be known
            let Some(elt) = dtype_size_bytes(&info.dtype) else {
                return Err(anyhow::anyhow!("unknown dtype: {}", info.dtype));
            };

            // offsets must be ordered and within buffer
            let (begin, end) = info.data_offsets;
            if begin > end {
                return Err(anyhow::anyhow!("invalid offsets (begin > end) for {k}"));
            }
            if end > data_len {
                return Err(anyhow::anyhow!("offset end out of bounds for {k}"));
            }

            // dtype*shape must match byte length, with overflow checks.
            let Some(numel) = checked_numel(&info.shape) else {
                return Err(anyhow::anyhow!("overflow computing numel for {k}"));
            };
            let Some(expected_bytes) = numel.checked_mul(elt) else {
                return Err(anyhow::anyhow!("overflow computing byte size for {k}"));
            };
            let actual_bytes = end.saturating_sub(begin);
            if expected_bytes != actual_bytes {
                return Err(anyhow::anyhow!(
                    "tensor byte size mismatch for {k}: expected {expected_bytes} got {actual_bytes}"
                ));
            }

            tensors.push((begin, end, k.clone(), info));
        }

        // Contiguity: 0..data_len must be fully covered without holes/overlaps.
        let mut cursor = 0usize;
        tensors.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.cmp(&b.1)).then(a.2.cmp(&b.2)));
        for (begin, end, name, _info) in tensors {
            if begin != cursor {
                return Err(anyhow::anyhow!(
                    "metadata does not fully cover buffer (hole or overlap): cursor={cursor} begin={begin} tensor={name}"
                ));
            }
            cursor = end;
        }
        if cursor != data_len {
            return Err(anyhow::anyhow!(
                "metadata does not fully cover buffer: end={cursor} data_len={data_len}"
            ));
        }

        Ok(())
    }

    /// Validate a safetensors file on disk (header + offsets).
    ///
    /// Reads the header only (bounded) and validates that offsets cover the data buffer.
    pub fn validate_file(path: &Path) -> Result<()> {
        let mut f = File::open(path)?;
        let file_len = f.metadata()?.len();
        if file_len < 8 {
            return Err(anyhow::anyhow!("file too small for safetensors header"));
        }

        let mut n_bytes = [0u8; 8];
        f.read_exact(&mut n_bytes)?;
        let n = u64::from_le_bytes(n_bytes);
        let n_usize = usize::try_from(n).map_err(|_| anyhow::anyhow!("header length overflow"))?;
        if n_usize == 0 {
            return Err(anyhow::anyhow!("header length is zero"));
        }
        if n_usize > MAX_HEADER_SIZE {
            return Err(anyhow::anyhow!("header too large"));
        }

        let data_start = 8u64
            .checked_add(n)
            .ok_or_else(|| anyhow::anyhow!("header length overflow"))?;
        if data_start > file_len {
            return Err(anyhow::anyhow!("invalid header length (past end of file)"));
        }

        let mut header = vec![0u8; n_usize];
        f.read_exact(&mut header)?;

        let data_len_u64 = file_len - data_start;
        let data_len =
            usize::try_from(data_len_u64).map_err(|_| anyhow::anyhow!("data length overflow"))?;
        validate_header_and_data_len(&header, data_len)
    }
}

/// Vector post-processing helpers (backed by `innr` SIMD primitives).
pub mod vector {
    use innr::{cosine, dot, norm};

    /// Threshold below which a vector is considered near-zero.
    const NORM_EPSILON: f32 = 1e-9;

    /// Compute the L2 norm of a vector.
    pub fn l2_norm(v: &[f32]) -> f32 {
        norm(v)
    }

    /// In-place L2 normalization.
    ///
    /// Returns the original norm. If the vector is near-zero (< `NORM_EPSILON`),
    /// this is a no-op and returns 0.0.
    pub fn l2_normalize_in_place(v: &mut [f32]) -> f32 {
        let n = norm(v);
        if n <= NORM_EPSILON {
            return 0.0;
        }
        let inv = 1.0 / n;
        for x in v.iter_mut() {
            *x *= inv;
        }
        n
    }

    /// Dot product.
    pub fn dot_f32(a: &[f32], b: &[f32]) -> f32 {
        dot(a, b)
    }

    /// Cosine similarity (handles zero vectors).
    pub fn cosine_f32(a: &[f32], b: &[f32]) -> f32 {
        cosine(a, b)
    }
}

#[cfg(feature = "candle-hf")]
mod candle_hf {
    use super::{EmbedMode, TextEmbedder};
    use anyhow::Result;
    use candle_core::{DType, Device, Module, Tensor};
    use candle_nn::VarBuilder;
    use hf_hub::api::sync::Api as HfApi;
    use once_cell::sync::OnceCell;
    use std::fs::File;
    use std::io::BufReader;
    use std::path::Path;
    use tokenizers::Tokenizer;

    // Architecture-specific model types from candle-transformers.
    use candle_transformers::models::bert::{BertModel, Config as BertConfig};
    use candle_transformers::models::distilbert::{Config as DistilBertConfig, DistilBertModel};
    use candle_transformers::models::jina_bert::{
        BertModel as JinaBertModel, Config as JinaBertConfig,
    };
    use candle_transformers::models::modernbert::{Config as ModernBertConfig, ModernBert};
    use candle_transformers::models::xlm_roberta::{Config as XlmRobertaConfig, XLMRobertaModel};
    // Stella v5 is not auto-detectable via model_type ("new" for 400M, "qwen2" for 1.5B).
    // It also needs two VarBuilders (base + embed head from separate safetensors files).
    // Supported via dedicated StellaEmbedder struct, not the generic LocalHfEmbedder path.
    use candle_transformers::models::stella_en_v5::{
        Config as StellaConfig, EmbedDim, EmbeddingModel as StellaEmbeddingModel,
    };

    static EMBEDDER: OnceCell<std::sync::Mutex<LocalHfEmbedder>> = OnceCell::new();

    /// Select the best available compute device for the enabled features, falling
    /// back to CPU. Metal on Apple Silicon (`metal` feature), CUDA on NVIDIA
    /// (`cuda` feature), else CPU. A failed accelerator init logs and falls back,
    /// so a missing GPU never breaks embedding.
    fn pick_device() -> Device {
        #[cfg(feature = "metal")]
        {
            match Device::new_metal(0) {
                Ok(d) => return d,
                Err(e) => eprintln!("embedd: Metal device unavailable, using CPU: {e}"),
            }
        }
        #[cfg(feature = "cuda")]
        {
            match Device::new_cuda(0) {
                Ok(d) => return d,
                Err(e) => eprintln!("embedd: CUDA device unavailable, using CPU: {e}"),
            }
        }
        Device::Cpu
    }

    /// Detected model architecture, used to dispatch the correct forward pass.
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum ModelArch {
        Bert,
        JinaBert,
        DistilBert,
        XlmRoberta,
        ModernBert,
    }

    /// Internal enum holding the loaded model variant.
    enum CandleModel {
        Bert(BertModel),
        JinaBert(JinaBertModel),
        DistilBert(DistilBertModel),
        XlmRoberta(XLMRobertaModel),
        ModernBert(ModernBert),
    }

    impl CandleModel {
        /// Run the forward pass appropriate for this architecture.
        ///
        /// All variants return `[bsz, seq, hidden]`.
        fn forward(
            &self,
            input_ids: &Tensor,
            token_type_ids: &Tensor,
            attention_mask: &Tensor,
        ) -> Result<Tensor> {
            match self {
                CandleModel::Bert(m) => {
                    Ok(m.forward(input_ids, token_type_ids, Some(attention_mask))?)
                }
                CandleModel::JinaBert(m) => Ok(m.forward(input_ids)?),
                CandleModel::DistilBert(m) => Ok(m.forward(input_ids, attention_mask)?),
                CandleModel::XlmRoberta(m) => {
                    Ok(m.forward(input_ids, attention_mask, token_type_ids, None, None, None)?)
                }
                CandleModel::ModernBert(m) => Ok(m.forward(input_ids, attention_mask)?),
            }
        }
    }

    /// Detect architecture from the raw config.json.
    fn detect_arch(config_json: &serde_json::Value) -> ModelArch {
        match config_json.get("model_type").and_then(|v| v.as_str()) {
            Some("distilbert") => ModelArch::DistilBert,
            Some("xlm-roberta") => ModelArch::XlmRoberta,
            Some("modernbert") => ModelArch::ModernBert,
            Some("bert") => {
                // JinaBERT uses ALiBi position embeddings.
                // See tests/arch_detection.rs for config.json test vectors.
                if config_json
                    .get("position_embedding_type")
                    .and_then(|v| v.as_str())
                    == Some("alibi")
                {
                    ModelArch::JinaBert
                } else {
                    ModelArch::Bert
                }
            }
            // Default to BERT for unknown types (best compatibility).
            _ => ModelArch::Bert,
        }
    }

    fn hidden_size(config_json: &serde_json::Value) -> usize {
        config_json
            .get("hidden_size")
            .or_else(|| config_json.get("dim"))
            .and_then(|v| v.as_u64())
            .unwrap_or(768) as usize
    }

    fn load_model(
        config_json: &serde_json::Value,
        vb: VarBuilder,
        arch: ModelArch,
    ) -> Result<CandleModel> {
        match arch {
            ModelArch::Bert => {
                let cfg: BertConfig = serde_json::from_value(config_json.clone())?;
                Ok(CandleModel::Bert(BertModel::load(vb, &cfg)?))
            }
            ModelArch::JinaBert => {
                let cfg: JinaBertConfig = serde_json::from_value(config_json.clone())?;
                Ok(CandleModel::JinaBert(JinaBertModel::new(vb, &cfg)?))
            }
            ModelArch::DistilBert => {
                let cfg: DistilBertConfig = serde_json::from_value(config_json.clone())?;
                Ok(CandleModel::DistilBert(DistilBertModel::load(vb, &cfg)?))
            }
            ModelArch::XlmRoberta => {
                let cfg: XlmRobertaConfig = serde_json::from_value(config_json.clone())?;
                Ok(CandleModel::XlmRoberta(XLMRobertaModel::new(&cfg, vb)?))
            }
            ModelArch::ModernBert => {
                let cfg: ModernBertConfig = serde_json::from_value(config_json.clone())?;
                Ok(CandleModel::ModernBert(ModernBert::load(vb, &cfg)?))
            }
        }
    }

    /// Sentence-embedding pooling strategy. BGE-family models pool the CLS token;
    /// most sentence-transformers (MiniLM, E5, GTE) mean-pool. Using the wrong one
    /// silently degrades retrieval quality, so it is read per-model from the
    /// `1_Pooling/config.json` the sentence-transformers convention ships.
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    enum Pooling {
        Cls,
        Mean,
    }

    /// Read the pooling strategy from a `1_Pooling/config.json` value; default to
    /// mean when the file is absent or does not request CLS pooling.
    fn pooling_from_config(v: &serde_json::Value) -> Pooling {
        if v.get("pooling_mode_cls_token")
            .and_then(|b| b.as_bool())
            .unwrap_or(false)
        {
            Pooling::Cls
        } else {
            Pooling::Mean
        }
    }

    /// Local embedding inference via HuggingFace Hub + Candle (CPU).
    ///
    /// Supports BERT, JinaBERT, DistilBERT, XLM-RoBERTa, and ModernBERT.
    /// Architecture is auto-detected from `config.json` `model_type` field.
    pub struct LocalHfEmbedder {
        pub model_id: String,
        tokenizer: Tokenizer,
        model: CandleModel,
        arch: ModelArch,
        device: Device,
        hidden_dim: usize,
        pooling: Pooling,
        query_prefix: String,
        doc_prefix: String,
        max_len: usize,
    }

    impl LocalHfEmbedder {
        /// Load from a local directory containing `config.json`, `tokenizer.json`,
        /// `model.safetensors`.
        pub fn from_dir(label: &str, dir: &Path) -> Result<Self> {
            let config_path = dir.join("config.json");
            let tok_path = dir.join("tokenizer.json");
            let weights_path = dir.join("model.safetensors");

            super::safetensors::validate_file(&weights_path)?;

            let config_json: serde_json::Value =
                serde_json::from_reader(BufReader::new(File::open(&config_path)?))?;
            let arch = detect_arch(&config_json);
            let hdim = hidden_size(&config_json);

            let device = pick_device();
            // SAFETY: safetensors header+offsets validated above. The mmap lifetime is
            // bound to VarBuilder which owns the mapping; no aliased mutation occurs.
            let vb = unsafe {
                VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, &device)?
            };
            let model = load_model(&config_json, vb, arch)?;

            let tokenizer = Tokenizer::from_file(tok_path).map_err(|e| anyhow::anyhow!("{e}"))?;

            let max_len = Self::resolve_max_len();
            let prompt = super::PromptTemplate::from_env_any();
            let pooling = File::open(dir.join("1_Pooling/config.json"))
                .ok()
                .and_then(|f| {
                    serde_json::from_reader::<_, serde_json::Value>(BufReader::new(f)).ok()
                })
                .map(|v| pooling_from_config(&v))
                .unwrap_or(Pooling::Mean);

            Ok(Self {
                model_id: label.to_string(),
                tokenizer,
                model,
                arch,
                device,
                hidden_dim: hdim,
                pooling,
                query_prefix: prompt.query_prefix,
                doc_prefix: prompt.doc_prefix,
                max_len,
            })
        }

        /// Load from a `ModelSource` (local dir or HuggingFace Hub id).
        pub fn new(model_source: &super::ModelSource) -> Result<Self> {
            match model_source {
                super::ModelSource::LocalDir(dir) => Self::from_dir(&dir.to_string_lossy(), dir),
                super::ModelSource::HuggingFaceModelId(model_id) => {
                    let api = HfApi::new()?;
                    let repo = api.model(model_id.to_string());
                    let config_path = repo.get("config.json")?;
                    let tok_path = repo.get("tokenizer.json")?;
                    let weights_path = repo.get("model.safetensors")?;

                    super::safetensors::validate_file(&weights_path)?;

                    let config_json: serde_json::Value =
                        serde_json::from_reader(BufReader::new(File::open(&config_path)?))?;
                    let arch = detect_arch(&config_json);
                    let hdim = hidden_size(&config_json);

                    let device = pick_device();
                    // SAFETY: see from_dir -- safetensors validated, mmap owned by VarBuilder.
                    let vb = unsafe {
                        VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, &device)?
                    };
                    let model = load_model(&config_json, vb, arch)?;

                    let tokenizer =
                        Tokenizer::from_file(tok_path).map_err(|e| anyhow::anyhow!("{e}"))?;

                    let max_len = Self::resolve_max_len();
                    let prompt = super::PromptTemplate::from_env_any();
                    let pooling = repo
                        .get("1_Pooling/config.json")
                        .ok()
                        .and_then(|p| File::open(p).ok())
                        .and_then(|f| {
                            serde_json::from_reader::<_, serde_json::Value>(BufReader::new(f)).ok()
                        })
                        .map(|v| pooling_from_config(&v))
                        .unwrap_or(Pooling::Mean);

                    Ok(Self {
                        model_id: model_id.to_string(),
                        tokenizer,
                        model,
                        arch,
                        device,
                        hidden_dim: hdim,
                        pooling,
                        query_prefix: prompt.query_prefix,
                        doc_prefix: prompt.doc_prefix,
                        max_len,
                    })
                }
            }
        }

        /// The detected model architecture.
        pub fn arch(&self) -> ModelArch {
            self.arch
        }

        /// Override the maximum token length (clamped to 2048).
        pub fn with_max_len(mut self, max_len: usize) -> Self {
            self.max_len = max_len.min(2048);
            self
        }

        /// Override the prompt template.
        pub fn with_prompt(mut self, prompt: super::PromptTemplate) -> Self {
            self.query_prefix = prompt.query_prefix;
            self.doc_prefix = prompt.doc_prefix;
            self
        }

        fn resolve_max_len() -> usize {
            std::env::var("EMBEDD_MAX_LEN")
                .or_else(|_| std::env::var("IKSH_EMBED_MAX_LEN"))
                .ok()
                .and_then(|s| s.parse::<usize>().ok())
                .unwrap_or(384)
                .min(2048)
        }

        /// Singleton access (reuses model weights across calls).
        ///
        /// Uses `ModelSource::from_env_any()` to resolve the model.
        pub fn get() -> Result<std::sync::MutexGuard<'static, LocalHfEmbedder>> {
            let m = EMBEDDER.get_or_try_init(|| {
                Ok::<std::sync::Mutex<LocalHfEmbedder>, anyhow::Error>(std::sync::Mutex::new(
                    Self::new(&super::ModelSource::from_env_any())?,
                ))
            })?;
            Ok(m.lock().expect("embedder mutex poisoned"))
        }

        /// Tokenize, pad, and run the forward pass. Returns the raw hidden states
        /// `[bsz, seq, hidden]`, the attention mask `[bsz, seq]`, and per-input
        /// real token counts (before padding).
        fn forward_inner(
            &self,
            texts: &[String],
            is_query: bool,
        ) -> Result<(Tensor, Tensor, Vec<usize>)> {
            let mut toks = Vec::with_capacity(texts.len());
            let mut masks = Vec::with_capacity(texts.len());
            let mut real_lens = Vec::with_capacity(texts.len());
            let mut max_seq = 0usize;

            for t in texts {
                let prefix = if is_query {
                    &self.query_prefix
                } else {
                    &self.doc_prefix
                };
                let s = format!("{prefix}{t}");
                let enc = self
                    .tokenizer
                    .encode(s, true)
                    .map_err(|e| anyhow::anyhow!("{e}"))?;
                let ids = enc.get_ids();
                let mask = enc.get_attention_mask();
                let len = ids.len().min(self.max_len);
                let ids: Vec<u32> = ids[..len].to_vec();
                let mask: Vec<u32> = mask[..len].to_vec();
                real_lens.push(len);
                max_seq = max_seq.max(ids.len());
                toks.push(ids);
                masks.push(mask);
            }

            let pad_id = self.tokenizer.token_to_id("[PAD]").unwrap_or(0);
            for (ids, mask) in toks.iter_mut().zip(masks.iter_mut()) {
                while ids.len() < max_seq {
                    ids.push(pad_id);
                    mask.push(0);
                }
            }

            let bsz = toks.len();
            let total = bsz * max_seq;
            let mut flat_ids = Vec::with_capacity(total);
            let mut flat_mask = Vec::with_capacity(total);
            for (ids, mask) in toks.into_iter().zip(masks) {
                flat_ids.extend_from_slice(&ids);
                flat_mask.extend_from_slice(&mask);
            }
            let input_ids =
                Tensor::from_vec(flat_ids, (bsz, max_seq), &self.device)?.to_dtype(DType::I64)?;
            let attention_mask =
                Tensor::from_vec(flat_mask, (bsz, max_seq), &self.device)?.to_dtype(DType::I64)?;
            let token_type_ids = Tensor::zeros((bsz, max_seq), DType::I64, &self.device)?;

            let ys = self
                .model
                .forward(&input_ids, &token_type_ids, &attention_mask)?;
            Ok((ys, attention_mask, real_lens))
        }

        fn embed_texts_inner(&self, texts: &[String], is_query: bool) -> Result<Vec<Vec<f32>>> {
            let (ys, attention_mask, _) = self.forward_inner(texts, is_query)?;

            // Pool per the model's convention, then L2 normalize.
            let ys = ys.to_dtype(DType::F32)?;
            let pooled = match self.pooling {
                // CLS pooling: the first token's hidden state (BGE family).
                Pooling::Cls => ys.narrow(1, 0, 1)?.squeeze(1)?,
                // Mean pooling over non-padding tokens (MiniLM, E5, GTE, ...).
                Pooling::Mean => {
                    let mask_f = attention_mask.to_dtype(DType::F32)?.unsqueeze(2)?;
                    let masked = ys.broadcast_mul(&mask_f)?;
                    let sum = masked.sum(1)?;
                    let denom = mask_f.sum(1)?.clamp(1e-6, f32::MAX)?;
                    sum.broadcast_div(&denom)?
                }
            };

            let norms = pooled
                .sqr()?
                .sum(1)?
                .sqrt()?
                .clamp(1e-12, f32::MAX)?
                .unsqueeze(1)?;
            let normed = pooled.broadcast_div(&norms)?;

            Ok(normed.to_vec2()?)
        }

        fn embed_tokens_inner(
            &self,
            texts: &[String],
            is_query: bool,
        ) -> Result<Vec<Vec<Vec<f32>>>> {
            let (ys, _attention_mask, real_lens) = self.forward_inner(texts, is_query)?;
            let ys = ys.to_dtype(DType::F32)?;

            // Extract per-token vectors, stripping padding.
            let all: Vec<Vec<Vec<f32>>> = ys.to_vec3()?;
            Ok(all
                .into_iter()
                .zip(real_lens)
                .map(|(token_vecs, real_len)| token_vecs[..real_len].to_vec())
                .collect())
        }
    }

    impl TextEmbedder for LocalHfEmbedder {
        fn embed_texts(&self, texts: &[String], mode: EmbedMode) -> Result<Vec<Vec<f32>>> {
            self.embed_texts_inner(texts, matches!(mode, EmbedMode::Query))
        }

        fn model_id(&self) -> Option<&str> {
            Some(&self.model_id)
        }

        fn dimension(&self) -> Option<usize> {
            Some(self.hidden_dim)
        }

        fn capabilities(&self) -> super::TextEmbedderCapabilities {
            super::TextEmbedderCapabilities {
                uses_embed_mode: super::PromptApplication::Internal,
                normalization: super::Normalization::L2Normalized,
                truncation: super::TruncationPolicy::Truncate {
                    max_len: Some(self.max_len),
                    direction: super::TruncationDirection::Right,
                },
            }
        }
    }

    impl super::TokenEmbedder for LocalHfEmbedder {
        fn embed_tokens(&self, texts: &[String], mode: EmbedMode) -> Result<Vec<Vec<Vec<f32>>>> {
            self.embed_tokens_inner(texts, matches!(mode, EmbedMode::Query))
        }
    }

    /// Stella v5 embedder with dedicated loading for the multi-file weight layout.
    ///
    /// Stella models store base transformer weights and embedding head weights in
    /// separate safetensors files (SentenceTransformers layout):
    /// - `model.safetensors` -- base transformer
    /// - `2_Dense_{dim}/model.safetensors` -- embedding head for output dimension
    ///
    /// The forward pass requires `&mut self` (candle KV cache), so this wraps in
    /// a Mutex internally for thread safety.
    pub struct StellaEmbedder {
        model: std::sync::Mutex<StellaEmbeddingModel>,
        model_id: String,
        embed_dim: usize,
        device: Device,
        tokenizer: Tokenizer,
        max_len: usize,
    }

    /// Map a numeric dim to the candle EmbedDim enum.
    fn to_embed_dim(dim: usize) -> Result<EmbedDim> {
        match dim {
            256 => Ok(EmbedDim::Dim256),
            768 => Ok(EmbedDim::Dim768),
            1024 => Ok(EmbedDim::Dim1024),
            2048 => Ok(EmbedDim::Dim2048),
            4096 => Ok(EmbedDim::Dim4096),
            6144 => Ok(EmbedDim::Dim6144),
            8192 => Ok(EmbedDim::Dim8192),
            _ => Err(anyhow::anyhow!(
                "unsupported Stella embed dim {dim}; supported: 256, 768, 1024, 2048, 4096, 6144, 8192"
            )),
        }
    }

    /// Detect Stella variant from config.json and build the candle Config.
    ///
    /// Stella 400M: hidden_size=1024, model_type="new"
    /// Stella 1.5B: hidden_size=1536, model_type="qwen2"
    fn stella_config(config_json: &serde_json::Value, embed_dim: EmbedDim) -> Result<StellaConfig> {
        let hidden = config_json
            .get("hidden_size")
            .and_then(|v| v.as_u64())
            .unwrap_or(0);
        match hidden {
            1024 => Ok(StellaConfig::new_400_m_v5(embed_dim)),
            1536 => Ok(StellaConfig::new_1_5_b_v5(embed_dim)),
            _ => Err(anyhow::anyhow!(
                "cannot detect Stella variant from hidden_size={hidden} (expected 1024 for 400M or 1536 for 1.5B)"
            )),
        }
    }

    impl StellaEmbedder {
        /// Load from a local directory with the SentenceTransformers layout.
        ///
        /// `embed_dim` selects which `2_Dense_{dim}/` subdirectory to load the
        /// embedding head from. Supported: 256, 768, 1024, 2048, 4096, 6144, 8192.
        pub fn from_dir(label: &str, dir: &Path, embed_dim: usize) -> Result<Self> {
            let edim = to_embed_dim(embed_dim)?;
            let config_path = dir.join("config.json");
            let tok_path = dir.join("tokenizer.json");
            let base_weights = dir.join("model.safetensors");
            let head_dir = dir.join(format!("2_Dense_{embed_dim}"));
            let head_weights = head_dir.join("model.safetensors");

            if !head_weights.exists() {
                return Err(anyhow::anyhow!(
                    "embed head not found: {} (available dims are in 2_Dense_*/)",
                    head_weights.display()
                ));
            }

            super::safetensors::validate_file(&base_weights)?;
            super::safetensors::validate_file(&head_weights)?;

            let config_json: serde_json::Value =
                serde_json::from_reader(BufReader::new(File::open(&config_path)?))?;
            let cfg = stella_config(&config_json, edim)?;

            let device = pick_device();

            // SAFETY: safetensors validated above, mmap owned by VarBuilder.
            let base_vb = unsafe {
                VarBuilder::from_mmaped_safetensors(&[base_weights], DType::F32, &device)?
            };
            let embed_vb = unsafe {
                VarBuilder::from_mmaped_safetensors(&[head_weights], DType::F32, &device)?
            };

            let model = StellaEmbeddingModel::new(&cfg, base_vb, embed_vb)?;
            let tokenizer = Tokenizer::from_file(tok_path).map_err(|e| anyhow::anyhow!("{e}"))?;

            Ok(Self {
                model: std::sync::Mutex::new(model),
                model_id: label.to_string(),
                embed_dim,
                device,
                tokenizer,
                max_len: 512,
            })
        }

        /// Load from HuggingFace Hub.
        pub fn from_hub(model_id: &str, embed_dim: usize) -> Result<Self> {
            let edim = to_embed_dim(embed_dim)?;
            let api = HfApi::new()?;
            let repo = api.model(model_id.to_string());

            let config_path = repo.get("config.json")?;
            let tok_path = repo.get("tokenizer.json")?;
            let base_weights = repo.get("model.safetensors")?;
            let head_path = format!("2_Dense_{embed_dim}/model.safetensors");
            let head_weights = repo.get(&head_path).map_err(|e| {
                anyhow::anyhow!("embed head {head_path} not found in {model_id}: {e}")
            })?;

            super::safetensors::validate_file(&base_weights)?;
            super::safetensors::validate_file(&head_weights)?;

            let config_json: serde_json::Value =
                serde_json::from_reader(BufReader::new(File::open(&config_path)?))?;
            let cfg = stella_config(&config_json, edim)?;

            let device = pick_device();

            // SAFETY: safetensors validated above, mmap owned by VarBuilder.
            let base_vb = unsafe {
                VarBuilder::from_mmaped_safetensors(&[base_weights], DType::F32, &device)?
            };
            let embed_vb = unsafe {
                VarBuilder::from_mmaped_safetensors(&[head_weights], DType::F32, &device)?
            };

            let model = StellaEmbeddingModel::new(&cfg, base_vb, embed_vb)?;
            let tokenizer = Tokenizer::from_file(tok_path).map_err(|e| anyhow::anyhow!("{e}"))?;

            Ok(Self {
                model: std::sync::Mutex::new(model),
                model_id: model_id.to_string(),
                embed_dim,
                device,
                tokenizer,
                max_len: 512,
            })
        }

        pub fn with_max_len(mut self, max_len: usize) -> Self {
            self.max_len = max_len.min(8192);
            self
        }
    }

    impl TextEmbedder for StellaEmbedder {
        fn embed_texts(&self, texts: &[String], _mode: EmbedMode) -> Result<Vec<Vec<f32>>> {
            // Tokenize.
            let mut toks = Vec::with_capacity(texts.len());
            let mut masks = Vec::with_capacity(texts.len());
            let mut max_seq = 0usize;

            for t in texts {
                let enc = self
                    .tokenizer
                    .encode(t.as_str(), true)
                    .map_err(|e| anyhow::anyhow!("{e}"))?;
                let ids = enc.get_ids();
                let mask = enc.get_attention_mask();
                let len = ids.len().min(self.max_len);
                toks.push(ids[..len].to_vec());
                masks.push(mask[..len].to_vec());
                max_seq = max_seq.max(len);
            }

            let pad_id = self.tokenizer.token_to_id("[PAD]").unwrap_or(0);
            for (ids, mask) in toks.iter_mut().zip(masks.iter_mut()) {
                while ids.len() < max_seq {
                    ids.push(pad_id);
                    mask.push(0);
                }
            }

            let bsz = toks.len();
            let total = bsz * max_seq;
            let mut flat_ids = Vec::with_capacity(total);
            let mut flat_mask = Vec::with_capacity(total);
            for (ids, mask) in toks.into_iter().zip(masks) {
                flat_ids.extend_from_slice(&ids);
                flat_mask.extend_from_slice(&mask);
            }
            let input_ids =
                Tensor::from_vec(flat_ids, (bsz, max_seq), &self.device)?.to_dtype(DType::I64)?;
            let attention_mask =
                Tensor::from_vec(flat_mask, (bsz, max_seq), &self.device)?.to_dtype(DType::I64)?;

            // Stella returns already-pooled [bsz, embed_dim] from forward.
            let mut guard = self.model.lock().expect("stella mutex poisoned");
            let pooled = guard.forward(&input_ids, &attention_mask)?;

            // L2 normalize.
            let norms = pooled
                .sqr()?
                .sum(1)?
                .sqrt()?
                .clamp(1e-12, f32::MAX)?
                .unsqueeze(1)?;
            let normed = pooled.broadcast_div(&norms)?;

            Ok(normed.to_vec2()?)
        }

        fn model_id(&self) -> Option<&str> {
            Some(&self.model_id)
        }

        fn dimension(&self) -> Option<usize> {
            Some(self.embed_dim)
        }

        fn capabilities(&self) -> super::TextEmbedderCapabilities {
            super::TextEmbedderCapabilities {
                uses_embed_mode: super::PromptApplication::None,
                normalization: super::Normalization::L2Normalized,
                truncation: super::TruncationPolicy::Truncate {
                    max_len: Some(self.max_len),
                    direction: super::TruncationDirection::Right,
                },
            }
        }
    }
}

#[cfg(feature = "candle-hf")]
pub use candle_hf::{LocalHfEmbedder, ModelArch, StellaEmbedder};

#[cfg(feature = "fastembed")]
pub use fastembed::FastembedReranker;

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

    #[test]
    fn prompt_template_apply_starts_with_prefix() {
        let p = PromptTemplate {
            query_prefix: "Q: ".into(),
            doc_prefix: "D: ".into(),
        };
        assert_eq!(p.apply(EmbedMode::Query, "x"), "Q: x");
        assert_eq!(p.apply(EmbedMode::Document, "x"), "D: x");
    }

    proptest! {
        #[test]
        fn prompt_applies_as_prefix(qp in ".*", dp in ".*", t in ".*") {
            let p = PromptTemplate { query_prefix: qp.clone(), doc_prefix: dp.clone() };

            let q = p.apply(EmbedMode::Query, &t);
            prop_assert!(q.starts_with(&qp));
            prop_assert!(q.ends_with(&t));

            let d = p.apply(EmbedMode::Document, &t);
            prop_assert!(d.starts_with(&dp));
            prop_assert!(d.ends_with(&t));
        }
    }

    proptest! {
        #[test]
        fn normalize_then_cosine_is_dot_for_nonzero(
            a in prop::collection::vec(-10.0f32..10.0f32, 1..64),
            b in prop::collection::vec(-10.0f32..10.0f32, 1..64),
        ) {
            // Match lengths.
            let n = a.len().min(b.len());
            let mut a = a[..n].to_vec();
            let mut b = b[..n].to_vec();

            let na = vector::l2_normalize_in_place(&mut a);
            let nb = vector::l2_normalize_in_place(&mut b);
            prop_assume!(na > 0.0 && nb > 0.0);

            let d = vector::dot_f32(&a, &b);
            let c = vector::cosine_f32(&a, &b);
            prop_assert!((d - c).abs() < 1e-3);
        }
    }

    #[test]
    #[allow(deprecated)]
    fn model_source_prefers_local_dir() {
        std::env::set_var("IKSH_EMBED_MODEL_DIR", "/tmp/somewhere");
        std::env::remove_var("IKSH_EMBED_MODEL");
        let src = ModelSource::from_iksh_env();
        assert!(matches!(src, ModelSource::LocalDir(_)));
        std::env::remove_var("IKSH_EMBED_MODEL_DIR");
    }

    #[test]
    fn safetensors_rejects_holes() {
        // One tensor covering 0..4, but data_len claims 8.
        let header = br#"{"t":{"dtype":"U8","shape":[4],"data_offsets":[0,4]}}"#;
        let err = safetensors::validate_header_and_data_len(header, 8).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("does not fully cover buffer"));
    }

    proptest! {
        #[test]
        fn safetensors_accepts_contiguous_layout(
            // Up to 16 tensors, each 0..128 bytes.
            sizes in prop::collection::vec(0usize..128, 1..16),
        ) {
            // Build a contiguous layout with dtype=U8 and shape=[size].
            let mut offset = 0usize;
            let mut entries = Vec::new();
            for (i, sz) in sizes.iter().enumerate() {
                let begin = offset;
                let end = offset + *sz;
                offset = end;
                entries.push(format!(
                    "\"t{i}\":{{\"dtype\":\"U8\",\"shape\":[{sz}],\"data_offsets\":[{begin},{end}]}}"
                ));
            }
            let json = format!("{{{}}}", entries.join(","));
            safetensors::validate_header_and_data_len(json.as_bytes(), offset).unwrap();
        }
    }

    // --- Caching tests ---

    /// A mock embedder that counts calls.
    struct CountingEmbedder {
        calls: std::sync::atomic::AtomicUsize,
    }

    impl CountingEmbedder {
        fn new() -> Self {
            Self {
                calls: std::sync::atomic::AtomicUsize::new(0),
            }
        }
        fn call_count(&self) -> usize {
            self.calls.load(std::sync::atomic::Ordering::SeqCst)
        }
    }

    impl TextEmbedder for CountingEmbedder {
        fn embed_texts(&self, texts: &[String], _mode: EmbedMode) -> anyhow::Result<Vec<Vec<f32>>> {
            self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            Ok(texts.iter().map(|t| vec![t.len() as f32]).collect())
        }
        fn model_id(&self) -> Option<&str> {
            Some("counter")
        }
    }

    #[test]
    fn caching_embedder_avoids_redundant_calls() {
        let inner = CountingEmbedder::new();
        let cached = CachingTextEmbedder::new(inner);

        let texts = vec!["hello".to_string(), "world".to_string()];
        let r1 = cached.embed_texts(&texts, EmbedMode::Document).unwrap();
        assert_eq!(cached.cache_len(), 2);
        assert_eq!(r1, vec![vec![5.0], vec![5.0]]);

        // Second call with same texts: should be fully cached.
        let r2 = cached.embed_texts(&texts, EmbedMode::Document).unwrap();
        assert_eq!(r1, r2);
        // Inner embedder called only once (for the first batch).
        assert_eq!(cached.inner.call_count(), 1);
    }

    #[test]
    fn caching_embedder_mode_separation() {
        let cached = CachingTextEmbedder::new(CountingEmbedder::new());
        let texts = vec!["hi".to_string()];

        cached.embed_texts(&texts, EmbedMode::Query).unwrap();
        cached.embed_texts(&texts, EmbedMode::Document).unwrap();
        // Different modes -> different cache entries.
        assert_eq!(cached.cache_len(), 2);
        assert_eq!(cached.inner.call_count(), 2);
    }

    #[test]
    fn caching_embedder_partial_hits() {
        let cached = CachingTextEmbedder::new(CountingEmbedder::new());

        // First: cache "a"
        cached
            .embed_texts(&["a".to_string()], EmbedMode::Document)
            .unwrap();
        assert_eq!(cached.inner.call_count(), 1);

        // Second: "a" (hit) + "b" (miss) -> only "b" forwarded
        let r = cached
            .embed_texts(&["a".to_string(), "b".to_string()], EmbedMode::Document)
            .unwrap();
        assert_eq!(r, vec![vec![1.0], vec![1.0]]);
        assert_eq!(cached.inner.call_count(), 2);
        assert_eq!(cached.cache_len(), 2);

        // Third: both cached
        cached
            .embed_texts(&["a".to_string(), "b".to_string()], EmbedMode::Document)
            .unwrap();
        assert_eq!(cached.inner.call_count(), 2); // no new calls
    }

    // --- Reranker tests ---

    /// A mock reranker that returns documents scored by their length (longer = higher score).
    struct LengthReranker;

    impl Reranker for LengthReranker {
        fn rerank(
            &self,
            _query: &str,
            documents: &[String],
            top_k: Option<usize>,
        ) -> anyhow::Result<Vec<RerankResult>> {
            let mut results: Vec<RerankResult> = documents
                .iter()
                .enumerate()
                .map(|(i, doc)| RerankResult {
                    index: i,
                    score: doc.len() as f32,
                })
                .collect();
            results.sort_by(|a, b| {
                b.score
                    .partial_cmp(&a.score)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
            if let Some(k) = top_k {
                results.truncate(k);
            }
            Ok(results)
        }

        fn model_id(&self) -> Option<&str> {
            Some("mock-length")
        }
    }

    #[test]
    fn rerank_result_sorted_by_score() {
        let docs = vec![
            "short".to_string(),
            "a]medium length".to_string(),
            "the longest document here".to_string(),
        ];
        let results = LengthReranker.rerank("query", &docs, None).unwrap();
        assert_eq!(results.len(), 3);
        // Longest first.
        assert_eq!(results[0].index, 2);
        assert_eq!(results[1].index, 1);
        assert_eq!(results[2].index, 0);
        // Scores descending.
        assert!(results[0].score >= results[1].score);
        assert!(results[1].score >= results[2].score);
    }

    #[test]
    fn rerank_top_k_truncation() {
        let docs = vec![
            "a".to_string(),
            "bb".to_string(),
            "ccc".to_string(),
            "dddd".to_string(),
        ];
        let results = LengthReranker.rerank("q", &docs, Some(2)).unwrap();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].index, 3);
        assert_eq!(results[1].index, 2);
    }

    #[test]
    fn batching_reranker_matches_unbatched() {
        let docs: Vec<String> = (0..10).map(|i| "x".repeat(i + 1)).collect();

        let unbatched = LengthReranker.rerank("q", &docs, None).unwrap();
        let batched = BatchingReranker::new(LengthReranker, 3)
            .rerank("q", &docs, None)
            .unwrap();

        assert_eq!(unbatched.len(), batched.len());
        // Same ordering and indices.
        for (u, b) in unbatched.iter().zip(batched.iter()) {
            assert_eq!(u.index, b.index);
            assert!((u.score - b.score).abs() < f32::EPSILON);
        }
    }

    #[test]
    fn batching_reranker_top_k() {
        let docs: Vec<String> = (0..10).map(|i| "x".repeat(i + 1)).collect();
        let results = BatchingReranker::new(LengthReranker, 4)
            .rerank("q", &docs, Some(3))
            .unwrap();
        assert_eq!(results.len(), 3);
        // Top 3 longest: indices 9, 8, 7.
        assert_eq!(results[0].index, 9);
        assert_eq!(results[1].index, 8);
        assert_eq!(results[2].index, 7);
    }

    #[test]
    fn batching_reranker_small_input_delegates() {
        // When docs fit in one batch, should behave identically.
        let docs = vec!["hello".to_string(), "world".to_string()];
        let direct = LengthReranker.rerank("q", &docs, None).unwrap();
        let batched = BatchingReranker::new(LengthReranker, 100)
            .rerank("q", &docs, None)
            .unwrap();
        assert_eq!(direct.len(), batched.len());
        for (d, b) in direct.iter().zip(batched.iter()) {
            assert_eq!(d.index, b.index);
        }
    }

    #[test]
    fn box_dyn_reranker() {
        let reranker: Box<dyn Reranker> = Box::new(LengthReranker);
        let docs = vec!["a".to_string(), "bbb".to_string()];
        let results = reranker.rerank("q", &docs, None).unwrap();
        assert_eq!(results[0].index, 1);
        assert_eq!(reranker.model_id(), Some("mock-length"));
    }
}