palace-rs 0.20.0

//! Text embedding using ONNX Runtime and the models in [`crate::embedding_models`].
//!
//! Each model loads lazily in its own `OnceCell` — commands that don't need
//! embeddings (e.g. `palace status`) start instantly, and a migration can hold
//! the active and the target model at once. Models are downloaded on first
//! use and cached in ~/.cache/huggingface.
//!
//! The *active* model is the one a palace's stored vectors were made with
//! (`meta.active_embedding_model`); [`crate::db::open`] sets it. Queries and
//! documents are embedded separately ([`embed_query`], [`embed_documents`])
//! because retrieval models such as BGE prefix queries with an instruction.

use anyhow::{Context, Result};
use ndarray::Array;
use once_cell::sync::OnceCell;
use ort::{
    session::{builder::GraphOptimizationLevel, Session},
    value::Value,
};
use std::collections::HashMap;
use std::sync::Mutex;
use std::thread::available_parallelism;
use std::time::Duration;
use unicode_normalization::UnicodeNormalization;

use crate::embedding_models::{ModelSpec, Pooling, MODELS};
use std::sync::atomic::{AtomicUsize, Ordering};

/// Embedding dimension of the default model (all-MiniLM-L6-v2). Other
/// models differ; use [`active_model`]`().dim` for the palace's real size.
pub const EMBEDDING_DIM: usize = 384;

// ── Minimal BERT WordPiece tokenizer ────────────────────────────────────────

struct BertTokenizer {
    token_to_id: HashMap<String, i64>,
    cls_id: i64,
    sep_id: i64,
    unk_id: i64,
    pad_id: i64,
}

impl BertTokenizer {
    fn from_vocab_file(path: &std::path::Path) -> Result<Self> {
        let content = std::fs::read_to_string(path).context("reading vocab.txt")?;
        let token_to_id: HashMap<String, i64> = content
            .lines()
            .enumerate()
            .map(|(i, tok)| (tok.to_string(), i as i64))
            .collect();

        let cls_id = *token_to_id
            .get("[CLS]")
            .ok_or_else(|| anyhow::anyhow!("vocab missing [CLS]"))?;
        let sep_id = *token_to_id
            .get("[SEP]")
            .ok_or_else(|| anyhow::anyhow!("vocab missing [SEP]"))?;
        let unk_id = *token_to_id
            .get("[UNK]")
            .ok_or_else(|| anyhow::anyhow!("vocab missing [UNK]"))?;
        let pad_id = *token_to_id.get("[PAD]").unwrap_or(&0);

        Ok(Self {
            token_to_id,
            cls_id,
            sep_id,
            unk_id,
            pad_id,
        })
    }

    fn encode(&self, text: &str, max_len: usize) -> EncodedInput {
        let cleaned = clean_and_lower(text);
        let basic_tokens = basic_tokenize(&cleaned);

        let mut wordpieces = Vec::new();
        for token in &basic_tokens {
            wordpiece_tokenize(token, &self.token_to_id, self.unk_id, &mut wordpieces);
        }

        // Reserve 2 for [CLS] and [SEP]
        let effective_max = max_len.saturating_sub(2);
        if wordpieces.len() > effective_max {
            wordpieces.truncate(effective_max);
        }

        let mut input_ids = Vec::with_capacity(wordpieces.len() + 2);
        input_ids.push(self.cls_id);
        input_ids.extend_from_slice(&wordpieces);
        input_ids.push(self.sep_id);

        let seq_len = input_ids.len();
        let attention_mask = vec![1i64; seq_len];
        let token_type_ids = vec![0i64; seq_len];

        EncodedInput {
            input_ids,
            attention_mask,
            token_type_ids,
        }
    }
}

struct EncodedInput {
    input_ids: Vec<i64>,
    attention_mask: Vec<i64>,
    token_type_ids: Vec<i64>,
}

fn clean_and_lower(text: &str) -> String {
    let mut out = String::with_capacity(text.len());
    for c in text.nfd() {
        if c.is_control() || c == '\0' {
            continue;
        }
        if is_combining_mark(c) {
            continue;
        }
        if is_whitespace_char(c) {
            out.push(' ');
        } else {
            for lc in c.to_lowercase() {
                out.push(lc);
            }
        }
    }
    out
}

fn is_combining_mark(c: char) -> bool {
    unicode_normalization::char::canonical_combining_class(c) != 0
}

fn is_whitespace_char(c: char) -> bool {
    matches!(c, ' ' | '\t' | '\n' | '\r') || {
        let cat = c as u32;
        // Unicode Zs category (space separators)
        cat == 0x00A0
            || cat == 0x1680
            || (0x2000..=0x200A).contains(&cat)
            || cat == 0x202F
            || cat == 0x205F
            || cat == 0x3000
    }
}

fn is_punctuation(c: char) -> bool {
    let cp = c as u32;
    (0x21..=0x2F).contains(&cp)
        || (0x3A..=0x40).contains(&cp)
        || (0x5B..=0x60).contains(&cp)
        || (0x7B..=0x7E).contains(&cp)
        || c.is_ascii_punctuation()
        || {
            // Unicode punctuation categories (P*)
            matches!(unicode_general_category(c), UnicodeCategory::Punctuation)
        }
}

fn is_cjk_character(c: char) -> bool {
    let cp = c as u32;
    (0x4E00..=0x9FFF).contains(&cp)
        || (0x3400..=0x4DBF).contains(&cp)
        || (0x20000..=0x2A6DF).contains(&cp)
        || (0x2A700..=0x2B73F).contains(&cp)
        || (0x2B740..=0x2B81F).contains(&cp)
        || (0x2B820..=0x2CEAF).contains(&cp)
        || (0xF900..=0xFAFF).contains(&cp)
        || (0x2F800..=0x2FA1F).contains(&cp)
}

#[derive(PartialEq)]
enum UnicodeCategory {
    Punctuation,
    Other,
}

fn unicode_general_category(c: char) -> UnicodeCategory {
    // Check Unicode general category for punctuation (Pc, Pd, Pe, Pf, Pi, Po, Ps)
    // Using char ranges from Unicode
    if c.is_ascii_punctuation() {
        return UnicodeCategory::Punctuation;
    }
    let cp = c as u32;
    // Common Unicode punctuation ranges
    if (0x2010..=0x2027).contains(&cp)
        || (0x2030..=0x205E).contains(&cp)
        || (0x2308..=0x230B).contains(&cp)
        || (0x2768..=0x2775).contains(&cp)
        || (0x27E6..=0x27EF).contains(&cp)
        || (0x2983..=0x2998).contains(&cp)
        || (0x3001..=0x3003).contains(&cp)
        || (0x3008..=0x3011).contains(&cp)
        || (0x3014..=0x301F).contains(&cp)
        || (0xFE30..=0xFE4F).contains(&cp)
        || (0xFF01..=0xFF0F).contains(&cp)
        || (0xFF1A..=0xFF20).contains(&cp)
        || (0xFF3B..=0xFF3D).contains(&cp)
        || (0xFF5B..=0xFF65).contains(&cp)
    {
        UnicodeCategory::Punctuation
    } else {
        UnicodeCategory::Other
    }
}

fn basic_tokenize(text: &str) -> Vec<String> {
    let mut tokens = Vec::new();
    for word in text.split_whitespace() {
        let mut current = String::new();
        for c in word.chars() {
            if is_punctuation(c) || is_cjk_character(c) {
                if !current.is_empty() {
                    tokens.push(std::mem::take(&mut current));
                }
                tokens.push(c.to_string());
            } else {
                current.push(c);
            }
        }
        if !current.is_empty() {
            tokens.push(current);
        }
    }
    tokens
}

fn wordpiece_tokenize(
    token: &str,
    vocab: &HashMap<String, i64>,
    unk_id: i64,
    output: &mut Vec<i64>,
) {
    if token.is_empty() {
        return;
    }
    if let Some(&id) = vocab.get(token) {
        output.push(id);
        return;
    }

    let chars: Vec<char> = token.chars().collect();
    let mut start = 0;
    let mut is_bad = false;

    while start < chars.len() {
        let mut end = chars.len();
        let mut found = false;

        while start < end {
            let substr: String = if start > 0 {
                let mut s = String::from("##");
                s.extend(&chars[start..end]);
                s
            } else {
                chars[start..end].iter().collect()
            };

            if let Some(&id) = vocab.get(&substr) {
                output.push(id);
                found = true;
                start = end;
                break;
            }
            end -= 1;
        }

        if !found {
            is_bad = true;
            break;
        }
    }

    if is_bad {
        // Clear any partial results and push UNK
        output.truncate(output.len().saturating_sub(chars.len()));
        output.push(unk_id);
    }
}

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

struct Embedder {
    spec: &'static ModelSpec,
    tokenizer: BertTokenizer,
    session: Mutex<Session>,
    need_token_type_ids: bool,
}

/// One lazily-built session per entry in [`MODELS`], by position.
static EMBEDDERS: once_cell::sync::Lazy<Vec<OnceCell<Embedder>>> =
    once_cell::sync::Lazy::new(|| MODELS.iter().map(|_| OnceCell::new()).collect());

/// Position in [`MODELS`] of the active model.
static ACTIVE: AtomicUsize = AtomicUsize::new(0);

/// How many times to attempt each model-file download before giving up.
const DOWNLOAD_ATTEMPTS: usize = 4;
/// Base delay for exponential backoff between download attempts.
const DOWNLOAD_BASE_DELAY: Duration = Duration::from_millis(500);

/// Retry `f` up to `attempts` times with exponential backoff, returning the
/// last error if every attempt fails.
///
/// This guards model downloads against transient HuggingFace failures —
/// notably HTTP 429 rate limiting, which otherwise fails CI runs that fetch the
/// model fresh.
fn retry<T, F>(attempts: usize, base_delay: Duration, mut f: F) -> Result<T>
where
    F: FnMut() -> Result<T>,
{
    let attempts = attempts.max(1);
    let mut last_err: Option<anyhow::Error> = None;
    for attempt in 0..attempts {
        match f() {
            Ok(value) => return Ok(value),
            Err(err) => last_err = Some(err),
        }
        if attempt + 1 < attempts {
            let factor = 2u32.saturating_pow(attempt as u32);
            let delay = base_delay.checked_mul(factor).unwrap_or(base_delay);
            if !delay.is_zero() {
                std::thread::sleep(delay);
            }
        }
    }
    Err(last_err.unwrap_or_else(|| anyhow::anyhow!("retry: no attempts were made")))
}

fn download_model_files(spec: &ModelSpec) -> Result<(std::path::PathBuf, std::path::PathBuf)> {
    use hf_hub::api::sync::Api;

    let api = Api::new().context("initializing HuggingFace API")?;
    let repo = api.model(spec.repo.to_string());

    let model_path = retry(DOWNLOAD_ATTEMPTS, DOWNLOAD_BASE_DELAY, || {
        repo.get(spec.model_file)
            .context("downloading ONNX model file")
    })?;
    let vocab_path = retry(DOWNLOAD_ATTEMPTS, DOWNLOAD_BASE_DELAY, || {
        repo.get(spec.vocab_file).context("downloading vocab file")
    })?;

    Ok((model_path, vocab_path))
}

/// ONNX intra-op threads: `PALACE_EMBED_THREADS` when set, else every core.
fn inference_threads() -> usize {
    std::env::var("PALACE_EMBED_THREADS")
        .ok()
        .and_then(|value| value.trim().parse::<usize>().ok())
        .filter(|&n| n > 0)
        .unwrap_or_else(|| available_parallelism().map(|n| n.get()).unwrap_or(1))
}

fn init_embedder(spec: &'static ModelSpec) -> Result<Embedder> {
    let (model_path, vocab_path) = download_model_files(spec)?;

    let session = Session::builder()?
        .with_optimization_level(GraphOptimizationLevel::Level3)?
        .with_intra_threads(inference_threads())?
        .commit_from_file(&model_path)
        .context("loading ONNX session")?;

    let need_token_type_ids = session
        .inputs()
        .iter()
        .any(|input| input.name() == "token_type_ids");

    let tokenizer =
        BertTokenizer::from_vocab_file(&vocab_path).context("loading BERT tokenizer")?;

    Ok(Embedder {
        spec,
        tokenizer,
        session: Mutex::new(session),
        need_token_type_ids,
    })
}

fn get_embedder(spec: &'static ModelSpec) -> Result<&'static Embedder> {
    EMBEDDERS[crate::embedding_models::index_of(spec)].get_or_try_init(|| {
        init_embedder(spec).with_context(|| format!("initializing embedder ({})", spec.id))
    })
}

/// Load `spec` now (downloading it if needed) without embedding anything.
pub fn load_model(spec: &'static ModelSpec) -> Result<()> {
    get_embedder(spec).map(|_| ())
}

/// The model this process embeds with by default.
pub fn active_model() -> &'static ModelSpec {
    &MODELS[ACTIVE.load(Ordering::Relaxed).min(MODELS.len() - 1)]
}

/// Make `id` the active model. Called when a palace is opened, from its
/// `meta.active_embedding_model`.
pub fn set_active_model(id: &str) -> Result<&'static ModelSpec> {
    let spec = crate::embedding_models::find(id).ok_or_else(|| {
        anyhow::anyhow!(
            "unknown embedding model '{id}' (known: {})",
            crate::embedding_models::ids()
        )
    })?;
    ACTIVE.store(crate::embedding_models::index_of(spec), Ordering::Relaxed);
    Ok(spec)
}

/// Warm the active embedder in a background thread so the model download
/// (23-90MB on a cold HuggingFace cache) and ONNX session build don't happen
/// inside a client's first `palace_mine` tool call, where they can exceed the
/// client's timeout on their own. Idempotent — `get_or_try_init` inside
/// `get_embedder` makes a redundant call a no-op.
pub fn warm_up_in_background() {
    let spec = active_model();
    let handle = std::thread::spawn(move || {
        if let Err(e) = get_embedder(spec) {
            tracing::warn!(error = %e, "Palace: background embedder warm-up failed");
        }
    });
    if let Ok(mut slot) = WARM_UP.lock() {
        *slot = Some(handle);
    }
}

/// The background warm-up thread, so shutdown can wait for it.
static WARM_UP: std::sync::Mutex<Option<std::thread::JoinHandle<()>>> = std::sync::Mutex::new(None);

/// Wait up to `limit` for a background warm-up to finish. Exiting the process
/// while ONNX Runtime is still building a session crashes it on teardown, so
/// short-lived processes call this before returning from `main`.
pub fn wait_for_warm_up(limit: Duration) {
    let Some(handle) = WARM_UP.lock().ok().and_then(|mut slot| slot.take()) else {
        return;
    };
    let deadline = std::time::Instant::now() + limit;
    while !handle.is_finished() {
        if std::time::Instant::now() >= deadline {
            tracing::warn!("Palace: exiting before the embedding model finished loading");
            return;
        }
        std::thread::sleep(Duration::from_millis(20));
    }
    let _ = handle.join();
}

/// Whether the active embedder has finished initializing. `mine()` uses this
/// to report a distinct "loading embedding model" progress phase instead of
/// going silent if a call arrives before warm-up completes.
pub fn is_ready() -> bool {
    EMBEDDERS[crate::embedding_models::index_of(active_model())]
        .get()
        .is_some()
}

/// Most texts sent to one ONNX call. Bigger batches stop paying off on CPU:
/// every row is padded to the longest one, so throughput falls as soon as a
/// batch mixes short and long texts.
const MAX_BATCH: usize = 8;

/// Embed `texts`, returning vectors in input order.
///
/// Texts are tokenized once, sorted by token length and run in sub-batches
/// of at most `MAX_BATCH`, so each batch pads only to neighbours of similar
/// length instead of to the longest text in the whole call.
fn run_inference(embedder: &Embedder, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
    if texts.is_empty() {
        return Ok(vec![]);
    }

    let encodings: Vec<EncodedInput> = texts
        .iter()
        .map(|text| embedder.tokenizer.encode(text, embedder.spec.max_seq_len))
        .collect();

    let mut order: Vec<usize> = (0..encodings.len()).collect();
    order.sort_by_key(|&i| encodings[i].input_ids.len());

    // int8 models quantize activations with scales taken over the whole
    // batch, so a text's vector would shift with its neighbours (cosine as
    // low as 0.976). One text per call keeps a text's vector the same
    // whether it is embedded alone or with others.
    let max_batch = if embedder.spec.precision == "int8" {
        1
    } else {
        MAX_BATCH
    };
    let mut results: Vec<Option<Vec<f32>>> = vec![None; encodings.len()];
    for chunk in order.chunks(max_batch) {
        let batch: Vec<&EncodedInput> = chunk.iter().map(|&i| &encodings[i]).collect();
        let vectors = run_batch(embedder, &batch)?;
        for (&i, vector) in chunk.iter().zip(vectors) {
            results[i] = Some(vector);
        }
    }
    results
        .into_iter()
        .map(|v| v.ok_or_else(|| anyhow::anyhow!("model returned too few embeddings")))
        .collect()
}

/// Run one padded batch through the model.
fn run_batch(embedder: &Embedder, encodings: &[&EncodedInput]) -> Result<Vec<Vec<f32>>> {
    let max_len = encodings
        .iter()
        .map(|e| e.input_ids.len())
        .max()
        .unwrap_or(0);
    let batch_size = encodings.len();
    let pad_id = embedder.tokenizer.pad_id;

    let mut ids_flat = Vec::with_capacity(batch_size * max_len);
    let mut mask_flat = Vec::with_capacity(batch_size * max_len);
    let mut type_ids_flat = Vec::with_capacity(batch_size * max_len);

    for enc in encodings {
        let seq_len = enc.input_ids.len();
        ids_flat.extend_from_slice(&enc.input_ids);
        mask_flat.extend_from_slice(&enc.attention_mask);
        type_ids_flat.extend_from_slice(&enc.token_type_ids);

        let padding = max_len.saturating_sub(seq_len);
        ids_flat.extend(std::iter::repeat_n(pad_id, padding));
        mask_flat.extend(std::iter::repeat_n(0i64, padding));
        type_ids_flat.extend(std::iter::repeat_n(0i64, padding));
    }

    let ids_array =
        Array::from_shape_vec((batch_size, max_len), ids_flat).context("shaping input_ids")?;
    let mask_array = Array::from_shape_vec((batch_size, max_len), mask_flat.clone())
        .context("shaping attention_mask")?;

    let mut session_inputs = ort::inputs![
        "input_ids" => Value::from_array(ids_array)?,
        "attention_mask" => Value::from_array(mask_array)?,
    ];

    if embedder.need_token_type_ids {
        let type_ids_array = Array::from_shape_vec((batch_size, max_len), type_ids_flat)
            .context("shaping token_type_ids")?;
        session_inputs.push((
            "token_type_ids".into(),
            Value::from_array(type_ids_array)?.into(),
        ));
    }

    let mut session = embedder
        .session
        .lock()
        .map_err(|e| anyhow::anyhow!("session lock poisoned: {e}"))?;
    let outputs = session
        .run(session_inputs)
        .context("running ONNX inference")?;

    let output_tensor = outputs
        .values()
        .next()
        .ok_or_else(|| anyhow::anyhow!("model returned no outputs"))?;

    let output_array = output_tensor
        .try_extract_array::<f32>()
        .context("extracting output tensor")?;

    let shape = output_array.shape();
    let embeddings = match shape.len() {
        2 => (0..batch_size)
            .map(|i| {
                let row: Vec<f32> = output_array
                    .slice(ndarray::s![i, ..])
                    .iter()
                    .copied()
                    .collect();
                l2_normalize(&row)
            })
            .collect(),
        3 => {
            let output_3d = output_array
                .into_dimensionality::<ndarray::Ix3>()
                .context("expected a 3-D output tensor")?;
            match embedder.spec.pooling {
                Pooling::Mean => mean_pool(output_3d.view(), &mask_flat, max_len),
                Pooling::Cls => cls_pool(output_3d.view()),
            }
        }
        _ => return Err(anyhow::anyhow!("unexpected output tensor shape: {shape:?}")),
    };

    Ok(embeddings)
}

/// Mean-pool token embeddings over the sequence dimension, masking out
/// padding, then L2-normalize each pooled vector. `mask_flat` is the
/// flattened `(batch_size, max_len)` attention mask used to build the model
/// input, indexed here the same way it was built.
///
/// Walks row slices (`sequence.outer_iter()`) rather than per-element
/// `output[[i, j, k]]` indexing — dynamic strided lookups, up to ~12.6M of
/// them for a full batch — and accumulates a single scalar `mask_sum` rather
/// than a `dim`-length vector recomputing the identical value (the mask
/// doesn't vary with the embedding dimension `k`) for every `k`.
fn mean_pool(output: ndarray::ArrayView3<f32>, mask_flat: &[i64], max_len: usize) -> Vec<Vec<f32>> {
    let batch_size = output.shape()[0];
    let dim = output.shape()[2];
    (0..batch_size)
        .map(|i| {
            let mut result = vec![0.0f32; dim];
            let mut mask_sum = 0.0f32;
            let sequence = output.index_axis(ndarray::Axis(0), i);
            for (j, row) in sequence.outer_iter().enumerate() {
                let m = mask_flat[i * max_len + j] as f32;
                if m == 0.0 {
                    continue;
                }
                mask_sum += m;
                for (acc, val) in result.iter_mut().zip(row.iter()) {
                    *acc += val * m;
                }
            }
            if mask_sum > 0.0 {
                for v in &mut result {
                    *v /= mask_sum;
                }
            }
            l2_normalize(&result)
        })
        .collect()
}

/// Take each sequence's first (`[CLS]`) token vector, L2-normalized.
fn cls_pool(output: ndarray::ArrayView3<f32>) -> Vec<Vec<f32>> {
    output
        .outer_iter()
        .map(|sequence| {
            let cls: Vec<f32> = sequence.row(0).iter().copied().collect();
            l2_normalize(&cls)
        })
        .collect()
}

fn l2_normalize(v: &[f32]) -> Vec<f32> {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    let epsilon = 1e-12;
    v.iter().map(|&x| x / (norm + epsilon)).collect()
}

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

/// Embed a search query with the active model.
pub fn embed_query(query: &str) -> Result<Vec<f32>> {
    embed_query_with(active_model(), query)
}

/// Embed a search query with `spec`, adding the model's query prefix.
pub fn embed_query_with(spec: &'static ModelSpec, query: &str) -> Result<Vec<f32>> {
    let embedder = get_embedder(spec)?;
    let text = format!("{}{query}", spec.query_prefix);
    run_inference(embedder, &[text.as_str()])
        .context("embedding query")?
        .pop()
        .ok_or_else(|| anyhow::anyhow!("embedding returned empty result"))
}

/// Embed one stored document (drawer, preference span, closet topic) with
/// the active model.
pub fn embed_document(text: &str) -> Result<Vec<f32>> {
    embed_documents(&[text])?
        .pop()
        .ok_or_else(|| anyhow::anyhow!("embedding returned empty result"))
}

/// Embed stored documents with the active model, in input order.
pub fn embed_documents(texts: &[&str]) -> Result<Vec<Vec<f32>>> {
    embed_documents_with(active_model(), texts)
}

/// Embed stored documents with `spec`, adding the model's document prefix.
pub fn embed_documents_with(spec: &'static ModelSpec, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
    if texts.is_empty() {
        return Ok(vec![]);
    }
    let embedder = get_embedder(spec)?;
    if spec.doc_prefix.is_empty() {
        return run_inference(embedder, texts).context("batch embedding");
    }
    let prefixed: Vec<String> = texts
        .iter()
        .map(|text| format!("{}{text}", spec.doc_prefix))
        .collect();
    let refs: Vec<&str> = prefixed.iter().map(String::as_str).collect();
    run_inference(embedder, &refs).context("batch embedding")
}

/// Embed a single piece of text with the active model.
#[deprecated(
    since = "0.18.0",
    note = "use embed_query for search text or embed_document for stored text"
)]
pub fn embed_one(text: &str) -> Result<Vec<f32>> {
    embed_document(text)
}

/// Embed multiple texts in one batch with the active model.
#[deprecated(since = "0.18.0", note = "use embed_documents")]
pub fn embed_batch(texts: &[&str]) -> Result<Vec<Vec<f32>>> {
    embed_documents(texts)
}

/// Serialize a f32 vector to little-endian bytes for SQLite BLOB storage.
pub fn vec_to_blob(v: &[f32]) -> Vec<u8> {
    v.iter().flat_map(|f| f.to_le_bytes()).collect()
}

/// Deserialize a BLOB back to a f32 vector.
pub fn blob_to_vec(blob: &[u8]) -> Vec<f32> {
    blob.chunks_exact(4)
        .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
        .collect()
}

/// Cosine similarity between two equal-length f32 vectors.
/// Returns a value in [-1.0, 1.0]; higher is more similar. Vectors of
/// different lengths come from different models and score 0.0.
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    if a.len() != b.len() {
        return 0.0;
    }
    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm_a == 0.0 || norm_b == 0.0 {
        return 0.0;
    }
    dot / (norm_a * norm_b)
}

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

    #[test]
    fn retry_recovers_from_transient_failures() {
        let calls = Cell::new(0);
        let result: Result<u32> = retry(4, Duration::ZERO, || {
            let n = calls.get() + 1;
            calls.set(n);
            if n < 3 {
                Err(anyhow::anyhow!("http status: 429"))
            } else {
                Ok(n)
            }
        });
        assert_eq!(result.unwrap(), 3, "should succeed once the call recovers");
        assert_eq!(
            calls.get(),
            3,
            "should stop retrying after the first success"
        );
    }

    #[test]
    fn retry_gives_up_after_exhausting_attempts() {
        let calls = Cell::new(0);
        let result: Result<u32> = retry(3, Duration::ZERO, || {
            calls.set(calls.get() + 1);
            Err(anyhow::anyhow!("persistent 429"))
        });
        assert!(result.is_err(), "should surface the failure after retries");
        assert_eq!(calls.get(), 3, "should attempt exactly `attempts` times");
    }

    /// Reference implementation of mean-pooling via the original
    /// per-element `output[[i, j, k]]` indexing, kept only in this test to
    /// cross-check `mean_pool`'s row-slice rewrite produces identical
    /// results.
    fn mean_pool_naive(
        output: &ndarray::Array3<f32>,
        mask_flat: &[i64],
        max_len: usize,
    ) -> Vec<Vec<f32>> {
        let batch_size = output.shape()[0];
        let dim = output.shape()[2];
        (0..batch_size)
            .map(|i| {
                let mut result = vec![0.0f32; dim];
                let mut mask_sum = vec![0.0f32; dim];
                let offset = i * max_len;
                for j in 0..max_len {
                    let m = mask_flat[offset + j] as f32;
                    for k in 0..dim {
                        let val = output[[i, j, k]];
                        result[k] += val * m;
                        mask_sum[k] += m;
                    }
                }
                for k in 0..dim {
                    if mask_sum[k] > 0.0 {
                        result[k] /= mask_sum[k];
                    }
                }
                l2_normalize(&result)
            })
            .collect()
    }

    #[test]
    fn mean_pool_ignores_padding_and_matches_hand_computed_result() {
        // batch=1, seq=3, dim=2. Position 2 is padding (mask=0) with
        // deliberately out-of-range values that must not affect the result.
        let output =
            ndarray::Array3::from_shape_vec((1, 3, 2), vec![1.0, 3.0, 2.0, 4.0, 99.0, 99.0])
                .unwrap();
        let mask_flat = [1i64, 1, 0];

        let pooled = mean_pool(output.view(), &mask_flat, 3);

        assert_eq!(pooled.len(), 1);
        // mean = [(1+2)/2, (3+4)/2] = [1.5, 3.5], ignoring the padded [99, 99].
        let expected = l2_normalize(&[1.5, 3.5]);
        for (actual, expected) in pooled[0].iter().zip(expected.iter()) {
            assert!(
                (actual - expected).abs() < 1e-6,
                "pooled={:?} expected={:?}",
                pooled[0],
                expected
            );
        }
    }

    #[test]
    fn mean_pool_matches_naive_per_element_reference_on_random_input() {
        // A fixed pseudo-random-looking tensor (no external RNG dependency)
        // across a batch/seq/dim shape and mask pattern large enough to
        // exercise multiple rows and a mix of masked/unmasked positions.
        let (batch, seq, dim) = (3, 5, 4);
        let mut values = Vec::with_capacity(batch * seq * dim);
        for i in 0..(batch * seq * dim) {
            let x = (i as f32) * 0.137 - 3.0;
            values.push(x.sin() * 2.0 + (i % 7) as f32 * 0.1);
        }
        let output = ndarray::Array3::from_shape_vec((batch, seq, dim), values).unwrap();
        let mask_flat: Vec<i64> = (0..batch * seq)
            .map(|i| if i % 4 == 3 { 0 } else { 1 })
            .collect();

        let fast = mean_pool(output.view(), &mask_flat, seq);
        let naive = mean_pool_naive(&output, &mask_flat, seq);

        assert_eq!(fast.len(), naive.len());
        for (row_fast, row_naive) in fast.iter().zip(naive.iter()) {
            for (a, b) in row_fast.iter().zip(row_naive.iter()) {
                assert!(
                    (a - b).abs() < 1e-6,
                    "fast={row_fast:?} naive={row_naive:?}"
                );
            }
        }
    }
}