//! CPU embedding backend using ndarray + system BLAS.
//!
//! Implements BERT inference on CPU using [`ndarray`] with system BLAS for
//! matrix operations. Weights are loaded directly from safetensors files
//! downloaded via `hf-hub`.
//!
//! Supports the `ClassicBert` family (BGE models): learned position embeddings,
//! GELU activation, and QKV projections with bias.

// Ensure the BLAS linker symbols are pulled in. On macOS we want
// Accelerate (the `cpu-accelerate` feature); elsewhere we want
// openblas (`cpu` feature). The blas-src crate is the shim.
#[cfg(all(target_os = "macos", feature = "cpu-accelerate"))]
extern crate accelerate_src;
#[cfg(feature = "cpu")]
extern crate blas_src;
#[cfg(all(not(target_os = "macos"), feature = "cpu"))]
extern crate openblas_src;

use std::f32::consts::PI;
use std::sync::Arc;

use hf_hub::api::sync::Api;
use ndarray::{Array1, Array2, ArrayView1, Axis, s};
use rayon::prelude::*;
use safetensors::SafeTensors;
use wide::f32x8;

use super::{DeviceHint, EmbedBackend, Encoding, RerankBackend};

// ---------------------------------------------------------------------------
// Architecture validation
// ---------------------------------------------------------------------------

/// Validate that the loaded weights are a recognized `ClassicBert` model.
///
/// Detect the architecture and tensor-name prefix.
///
/// `ClassicBert` (BGE-small embedding) uses bare names like
/// `embeddings.position_embeddings.weight`. HuggingFace `transformers`
/// exports (used by MS-MARCO cross-encoders) wrap with `bert.` prefix.
/// We try both and return whichever resolves.
///
/// Returns the prefix string (empty for bare, `"bert."` for HF
/// exports). The caller passes it to subsequent tensor lookups.
fn detect_variant(tensors: &SafeTensors<'_>) -> crate::Result<&'static str> {
    if tensors
        .tensor("embeddings.position_embeddings.weight")
        .is_ok()
    {
        Ok("")
    } else if tensors
        .tensor("bert.embeddings.position_embeddings.weight")
        .is_ok()
    {
        Ok("bert.")
    } else {
        Err(crate::Error::Other(anyhow::anyhow!(
            "unrecognized model architecture: no position_embeddings found"
        )))
    }
}

// ---------------------------------------------------------------------------
// BERT model configuration
// ---------------------------------------------------------------------------

/// Configuration for a BERT-style encoder model.
///
/// Matches the `config.json` schema from `HuggingFace` model repos.
#[derive(Debug, Clone)]
struct BertConfig {
    /// Hidden dimension (e.g. 384 for bge-small).
    hidden_size: i32,
    /// Number of transformer layers.
    num_hidden_layers: i32,
    /// Number of attention heads.
    num_attention_heads: i32,
    /// Maximum sequence length (512 for ClassicBert).
    max_position_embeddings: i32,
    /// Layer norm epsilon.
    layer_norm_eps: f32,
}

impl BertConfig {
    /// Parse from a `config.json` value.
    #[expect(
        clippy::cast_possible_truncation,
        reason = "config values are small ints/floats that fit in i32/f32"
    )]
    fn from_json(v: &serde_json::Value) -> crate::Result<Self> {
        let get_i32 = |key: &str| -> crate::Result<i32> {
            v.get(key)
                .and_then(serde_json::Value::as_i64)
                .map(|n| n as i32)
                .ok_or_else(|| crate::Error::Other(anyhow::anyhow!("missing config key: {key}")))
        };
        let get_f64 = |key: &str| -> crate::Result<f64> {
            v.get(key)
                .and_then(serde_json::Value::as_f64)
                .ok_or_else(|| crate::Error::Other(anyhow::anyhow!("missing config key: {key}")))
        };

        let layer_norm_eps =
            get_f64("layer_norm_epsilon").or_else(|_| get_f64("layer_norm_eps"))? as f32;

        Ok(Self {
            hidden_size: get_i32("hidden_size")?,
            num_hidden_layers: get_i32("num_hidden_layers")?,
            num_attention_heads: get_i32("num_attention_heads")?,
            max_position_embeddings: get_i32("max_position_embeddings").unwrap_or(512),
            layer_norm_eps,
        })
    }
}

// ---------------------------------------------------------------------------
// Safetensors -> ndarray helpers
// ---------------------------------------------------------------------------

/// Load a named tensor from safetensors as `Array2<f32>`.
///
/// The tensor must be stored in `f32` (little-endian) format. Returns an
/// error if the tensor is missing or the byte count does not match the shape.
fn load_tensor2(tensors: &SafeTensors<'_>, name: &str) -> crate::Result<Array2<f32>> {
    let tensor = tensors
        .tensor(name)
        .map_err(|_| crate::Error::Other(anyhow::anyhow!("missing weight: {name}")))?;
    let shape = tensor.shape();
    if shape.len() != 2 {
        return Err(crate::Error::Other(anyhow::anyhow!(
            "expected 2D tensor for {name}, got {}D",
            shape.len()
        )));
    }
    let data: Vec<f32> = tensor
        .data()
        .chunks_exact(4)
        .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
        .collect();
    Array2::from_shape_vec((shape[0], shape[1]), data)
        .map_err(|e| crate::Error::Other(anyhow::anyhow!("shape error for {name}: {e}")))
}

/// Load a named tensor from safetensors as `Array1<f32>`.
///
/// The tensor must be stored in `f32` (little-endian) format.
fn load_tensor1(tensors: &SafeTensors<'_>, name: &str) -> crate::Result<Array1<f32>> {
    let tensor = tensors
        .tensor(name)
        .map_err(|_| crate::Error::Other(anyhow::anyhow!("missing weight: {name}")))?;
    let shape = tensor.shape();
    if shape.len() != 1 {
        return Err(crate::Error::Other(anyhow::anyhow!(
            "expected 1D tensor for {name}, got {}D",
            shape.len()
        )));
    }
    let data: Vec<f32> = tensor
        .data()
        .chunks_exact(4)
        .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
        .collect();
    Ok(Array1::from_vec(data))
}

/// Optionally load a 1D tensor -- returns `None` if missing.
fn try_load_tensor1(tensors: &SafeTensors<'_>, name: &str) -> crate::Result<Option<Array1<f32>>> {
    if tensors.tensor(name).is_ok() {
        Ok(Some(load_tensor1(tensors, name)?))
    } else {
        Ok(None)
    }
}

// ---------------------------------------------------------------------------
// Layer norm helper
// ---------------------------------------------------------------------------

/// Apply layer normalization to a 1D hidden vector (single token).
///
/// Computes: `(x - mean) / sqrt(var + eps) * weight + bias`
///
/// Vectorized: mean and variance reductions both use `f32x8` accumulators
/// (one horizontal sum at the end). The fused normalize-and-affine pass
/// produces the output via a single `f32x8` write per 8-lane chunk
/// (FMA-friendly: `(x - mean) * inv_std * weight + bias`). Three scalar
/// `mapv` passes collapse into one vector pass; the per-token cost
/// on hidden_dim=384 drops to about a quarter.
fn layer_norm(
    x: &ArrayView1<'_, f32>,
    weight: &Array1<f32>,
    bias: &Array1<f32>,
    eps: f32,
) -> Array1<f32> {
    let n = x.len();
    debug_assert_eq!(weight.len(), n);
    debug_assert_eq!(bias.len(), n);

    let x_s = x.as_slice().expect("layer_norm: x must be contiguous");
    let w_s = weight
        .as_slice()
        .expect("layer_norm: weight must be contiguous");
    let b_s = bias
        .as_slice()
        .expect("layer_norm: bias must be contiguous");

    // Pass 1: vectorized mean = sum(x) / n.
    let sum: f32 = horizontal_sum_f32x8(x_s);
    #[expect(clippy::cast_precision_loss, reason = "hidden_dim is well under 2^24")]
    let mean = sum / n as f32;

    // Pass 2: vectorized variance = sum((x - mean)^2) / n.
    let mean_v = f32x8::splat(mean);
    let body = n - (n % 8);
    let (x_body, x_tail) = x_s.split_at(body);
    let mut var_acc = f32x8::splat(0.0);
    for chunk in x_body.chunks_exact(8) {
        let v = f32x8::from(<[f32; 8]>::try_from(chunk).unwrap());
        let diff = v - mean_v;
        var_acc = diff.mul_add(diff, var_acc);
    }
    let mut var_sum: f32 = var_acc.as_array().iter().sum();
    for &v in x_tail {
        let d = v - mean;
        var_sum += d * d;
    }
    #[expect(clippy::cast_precision_loss, reason = "hidden_dim is well under 2^24")]
    let var = var_sum / n as f32;
    let inv_std = 1.0 / (var + eps).sqrt();
    let inv_std_v = f32x8::splat(inv_std);

    // Pass 3: fused (x - mean) * inv_std * weight + bias, single write.
    // Pre-zero `out` so the body chunks_exact_mut can write in place;
    // the SIMD memset is far cheaper than the FMA work that follows.
    let mut out: Vec<f32> = vec![0.0; n];
    let (w_body, w_tail) = w_s.split_at(body);
    let (b_body, b_tail) = b_s.split_at(body);
    let (out_body, out_tail) = out.split_at_mut(body);
    for (((x_chunk, w_chunk), b_chunk), o_chunk) in x_body
        .chunks_exact(8)
        .zip(w_body.chunks_exact(8))
        .zip(b_body.chunks_exact(8))
        .zip(out_body.chunks_exact_mut(8))
    {
        let v = f32x8::from(<[f32; 8]>::try_from(x_chunk).unwrap());
        let w = f32x8::from(<[f32; 8]>::try_from(w_chunk).unwrap());
        let b = f32x8::from(<[f32; 8]>::try_from(b_chunk).unwrap());
        // (v - mean) * inv_std * weight + bias
        // Lowers to two FMAs on aarch64 (vfmaq_f32 / vfmsq_f32).
        let normed = (v - mean_v) * inv_std_v;
        let result = normed.mul_add(w, b);
        o_chunk.copy_from_slice(result.as_array());
    }
    for (((&xv, &wv), &bv), o) in x_tail
        .iter()
        .zip(w_tail.iter())
        .zip(b_tail.iter())
        .zip(out_tail.iter_mut())
    {
        *o = (xv - mean) * inv_std * wv + bv;
    }
    Array1::from(out)
}

/// Horizontal sum of a slice, vectorized.
fn horizontal_sum_f32x8(v: &[f32]) -> f32 {
    let n = v.len();
    let body = n - (n % 8);
    let (body_slice, tail) = v.split_at(body);
    let mut acc = f32x8::splat(0.0);
    for chunk in body_slice.chunks_exact(8) {
        let x = f32x8::from(<[f32; 8]>::try_from(chunk).unwrap());
        acc += x;
    }
    let mut s: f32 = acc.as_array().iter().sum();
    for &x in tail {
        s += x;
    }
    s
}

// ---------------------------------------------------------------------------
// Activation helpers
// ---------------------------------------------------------------------------

/// GELU activation (tanh approximation), applied in place to an
/// `ndarray::Array2<f32>` over its underlying contiguous slice.
///
/// `gelu(x) = x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`
///
/// Vectorized via `wide::f32x8`. `wide` exposes `exp` but not `tanh`,
/// so we expand `tanh(z) = 1 - 2 / (exp(2z) + 1)` and call `exp` once
/// per 8-lane chunk. The FFN intermediate vector after the
/// hidden→4×hidden projection (1536 floats for 384-dim BERT, 3072 for
/// 768-dim ModernBERT) is the hottest place this fires; one rerank
/// query over 100 candidates × 12 layers × ~6 tokens × 1536 dims =
/// ~11M GELU evaluations. Replacing the scalar `tanhf` with one
/// vectorized `exp` per 8 elements is the largest single SIMD win
/// in the BERT inference path.
///
/// Operates on `Array2::as_slice_mut()` so caller's matrix must be
/// row-major contiguous (which it is, as produced by `ndarray::dot`).
fn gelu_inplace(intermediate: &mut Array2<f32>) {
    let slice = intermediate
        .as_slice_mut()
        .expect("gelu_inplace: matrix must be contiguous (row-major)");
    let n = slice.len();
    let body = n - (n % 8);
    let (body_slice, tail) = slice.split_at_mut(body);

    let one = f32x8::splat(1.0);
    let two = f32x8::splat(2.0);
    let coef = f32x8::splat((2.0_f32 / PI).sqrt());
    let cubic_c = f32x8::splat(0.044_715);

    for chunk in body_slice.chunks_exact_mut(8) {
        let x = f32x8::from(<[f32; 8]>::try_from(&*chunk).unwrap());
        let x_sq = x * x;
        let inner = (coef * x).mul_add(cubic_c.mul_add(x_sq, one), f32x8::splat(0.0));
        // inner = coef * (x + 0.044715 * x^3); above we rewrote as
        //   coef * x * (1 + 0.044715 * x^2)
        // which is equivalent but uses one fewer multiply and an FMA.

        // tanh(2*inner) -- but the GELU form takes tanh(inner) directly:
        //   gelu(x) = x * 0.5 * (1 + tanh(inner))
        // Expand tanh(inner) = 1 - 2 / (exp(2*inner) + 1):
        //   1 + tanh(inner) = 2 - 2 / (exp(2*inner) + 1)
        //                   = 2 * (1 - 1 / (exp(2*inner) + 1))
        // So gelu(x) = x * (1 - 1 / (exp(2*inner) + 1)).
        let e = (two * inner).exp();
        let frac = one / (e + one);
        let result = x * (one - frac);
        chunk.copy_from_slice(result.as_array());
    }
    for x in tail.iter_mut() {
        *x *= 0.5 * (1.0 + (coef.as_array()[0] * (*x + 0.044_715 * x.powi(3))).tanh());
    }
}

/// Softmax along the last axis of a 1D slice (in-place).
///
/// Uses the numerically stable `exp(x - max) / sum(exp(x - max))` form.
///
/// Vectorized: the max reduction, the `exp(x - max)` and accumulation
/// pass, and the final divide all use `f32x8`. The max reduction needs
/// a lane-wise max (`wide::f32x8` doesn't expose one directly); we
/// horizontally fold to scalar after a vector running-max accumulator.
///
/// Called per attention head per layer per token in BERT — ~10K
/// invocations per query at the operative shapes. The hot lane is
/// `(x - max).exp()` which becomes one vector `exp` per 8 elements
/// instead of one scalar `expf` per element.
fn softmax_inplace(vals: &mut [f32]) {
    let n = vals.len();
    if n == 0 {
        return;
    }

    // Pass 1: max reduction.
    let body = n - (n % 8);
    let (body_slice, tail) = vals.split_at_mut(body);
    let max = if body_slice.is_empty() {
        tail.iter().copied().fold(f32::NEG_INFINITY, f32::max)
    } else {
        // Vector running max -> horizontal fold -> tail fold.
        let mut acc = f32x8::from(<[f32; 8]>::try_from(&body_slice[..8]).unwrap());
        for chunk in body_slice[8..].chunks_exact(8) {
            let v = f32x8::from(<[f32; 8]>::try_from(chunk).unwrap());
            acc = acc.max(v);
        }
        let vec_max = acc
            .as_array()
            .iter()
            .copied()
            .fold(f32::NEG_INFINITY, f32::max);
        tail.iter().copied().fold(vec_max, f32::max)
    };
    let max_v = f32x8::splat(max);

    // Pass 2: in-place exp(x - max); accumulate sum.
    let mut sum_acc = f32x8::splat(0.0);
    for chunk in body_slice.chunks_exact_mut(8) {
        let x = f32x8::from(<[f32; 8]>::try_from(&*chunk).unwrap());
        let e = (x - max_v).exp();
        sum_acc += e;
        chunk.copy_from_slice(e.as_array());
    }
    let mut sum: f32 = sum_acc.as_array().iter().sum();
    for v in tail.iter_mut() {
        *v = (*v - max).exp();
        sum += *v;
    }

    // Pass 3: divide by sum.
    let inv_sum = 1.0 / sum;
    let inv_sum_v = f32x8::splat(inv_sum);
    for chunk in body_slice.chunks_exact_mut(8) {
        let x = f32x8::from(<[f32; 8]>::try_from(&*chunk).unwrap());
        chunk.copy_from_slice((x * inv_sum_v).as_array());
    }
    for v in tail.iter_mut() {
        *v *= inv_sum;
    }
}

// ---------------------------------------------------------------------------
// BERT embeddings layer
// ---------------------------------------------------------------------------

/// BERT embeddings layer: word + position + `token_type` + `LayerNorm`.
///
/// Position and `token_type` embeddings are summed with word embeddings.
#[derive(Debug)]
struct CpuBertEmbeddings {
    /// Word embedding table `[vocab_size, hidden]`.
    word_embeddings: Array2<f32>,
    /// Learned position embeddings (`ClassicBert` only) `[max_seq, hidden]`.
    position_embeddings: Option<Array2<f32>>,
    /// Token type embeddings (`ClassicBert` only) `[2, hidden]`.
    token_type_embeddings: Option<Array2<f32>>,
    /// Layer norm weight `[hidden]`.
    layer_norm_weight: Array1<f32>,
    /// Layer norm bias `[hidden]`.
    layer_norm_bias: Array1<f32>,
    /// Layer norm epsilon.
    layer_norm_eps: f32,
}

impl CpuBertEmbeddings {
    /// Forward pass: look up embeddings, sum, and normalize.
    ///
    /// Returns one `[seq, hidden]` matrix per batch item.
    #[expect(
        clippy::cast_sign_loss,
        clippy::cast_possible_truncation,
        reason = "token IDs from tokenizer are always non-negative and fit in usize"
    )]
    fn forward(&self, encodings: &[Encoding]) -> Vec<Array2<f32>> {
        let hidden = self.word_embeddings.shape()[1];

        encodings
            .iter()
            .map(|enc| {
                let seq_len = enc.input_ids.len();
                let mut output = Array2::<f32>::zeros((seq_len, hidden));

                for (t, &id) in enc.input_ids.iter().enumerate() {
                    let word_row = self.word_embeddings.row(id as usize);
                    output.row_mut(t).assign(&word_row);

                    if let Some(ref pos_emb) = self.position_embeddings {
                        let pos_row = pos_emb.row(t);
                        output.row_mut(t).zip_mut_with(&pos_row, |o, &p| *o += p);
                    }

                    if let Some(ref tok_emb) = self.token_type_embeddings {
                        let type_id = enc.token_type_ids[t] as usize;
                        let tok_row = tok_emb.row(type_id);
                        output.row_mut(t).zip_mut_with(&tok_row, |o, &p| *o += p);
                    }

                    let normed = layer_norm(
                        &output.row(t),
                        &self.layer_norm_weight,
                        &self.layer_norm_bias,
                        self.layer_norm_eps,
                    );
                    output.row_mut(t).assign(&normed);
                }

                output
            })
            .collect()
    }
}

// ---------------------------------------------------------------------------
// Self-attention
// ---------------------------------------------------------------------------

/// Self-attention sub-layer within a BERT encoder layer.
///
/// Uses a fused QKV projection: a single `[3*hidden, hidden]` weight matrix
/// produces Q, K, V in one matmul, then splits the result.
/// Projections include bias terms; no rotary encoding.
#[derive(Debug)]
struct CpuBertSelfAttention {
    /// Fused Q/K/V weight matrix `[3*hidden, hidden]`.
    qkv_weight: Array2<f32>,
    /// Fused Q/K/V bias `[3*hidden]`.
    qkv_bias: Option<Array1<f32>>,
    /// Output projection weight `[hidden, hidden]`.
    output_weight: Array2<f32>,
    /// Output projection bias `[hidden]`.
    output_bias: Option<Array1<f32>>,
    /// Post-attention `LayerNorm` weight `[hidden]`.
    output_ln_weight: Array1<f32>,
    /// Post-attention `LayerNorm` bias `[hidden]`.
    output_ln_bias: Array1<f32>,
    /// Number of attention heads.
    num_heads: i32,
    /// Dimension per head (`hidden / num_heads`).
    head_dim: i32,
    /// Layer norm epsilon.
    layer_norm_eps: f32,
}

impl CpuBertSelfAttention {
    /// Scaled dot-product multi-head attention with residual + `LayerNorm`.
    ///
    /// Both variants use post-norm: attention -> residual -> `LayerNorm`.
    /// `hidden` is `[seq, hidden]` (single batch item). `mask` is `[seq]`
    /// with 0.0 for real tokens and a large negative value for padding.
    #[expect(
        clippy::cast_sign_loss,
        clippy::cast_precision_loss,
        reason = "num_heads/head_dim are small positive ints from config"
    )]
    fn forward(&self, hidden: &Array2<f32>, mask: &Array1<f32>) -> crate::Result<Array2<f32>> {
        let seq = hidden.shape()[0];
        let nh = self.num_heads as usize;
        let hd = self.head_dim as usize;
        let hidden_dim = nh * hd;

        // --- Fused QKV projection: [seq, hidden] @ [hidden, 3*hidden] => [seq, 3*hidden] ---
        let qkv = hidden.dot(&self.qkv_weight.t());
        let qkv = if let Some(ref bias) = self.qkv_bias {
            qkv + bias
        } else {
            qkv
        };

        // Split into Q, K, V each [seq, hidden]
        let q = qkv.slice(s![.., 0..hidden_dim]).to_owned();
        let k = qkv.slice(s![.., hidden_dim..2 * hidden_dim]).to_owned();
        let v = qkv.slice(s![.., 2 * hidden_dim..3 * hidden_dim]).to_owned();

        // --- Per-head attention ---
        // We iterate over heads since ndarray doesn't have native 4D matmul.
        let mut context = Array2::<f32>::zeros((seq, hidden_dim));

        for h in 0..nh {
            let col_start = h * hd;
            let col_end = col_start + hd;

            // Extract head slices: [seq, head_dim]
            let q_h = q.slice(s![.., col_start..col_end]);
            let k_h = k.slice(s![.., col_start..col_end]);
            let v_h = v.slice(s![.., col_start..col_end]);

            // scores = Q @ K^T / sqrt(head_dim)  => [seq, seq]
            let scale = 1.0 / (hd as f32).sqrt();
            let mut scores = q_h.dot(&k_h.t());
            scores.mapv_inplace(|v| v * scale);

            // Add attention mask: broadcast [seq] to each row
            for mut row in scores.rows_mut() {
                row.zip_mut_with(mask, |s, &m| *s += m);
            }

            // Softmax along last axis (each row)
            for mut row in scores.rows_mut() {
                softmax_inplace(row.as_slice_mut().ok_or_else(|| {
                    crate::Error::Other(anyhow::anyhow!("attention scores not contiguous"))
                })?);
            }

            // context_h = scores @ V_h  => [seq, head_dim]
            let ctx_h = scores.dot(&v_h);
            context.slice_mut(s![.., col_start..col_end]).assign(&ctx_h);
        }

        // --- Output projection ---
        let projected = context.dot(&self.output_weight.t());
        let projected = if let Some(ref bias) = self.output_bias {
            projected + bias
        } else {
            projected
        };

        // --- Residual + LayerNorm ---
        let residual = hidden + &projected;
        let mut output = Array2::<f32>::zeros((seq, hidden_dim));
        for t in 0..seq {
            let normed = layer_norm(
                &residual.row(t),
                &self.output_ln_weight,
                &self.output_ln_bias,
                self.layer_norm_eps,
            );
            output.row_mut(t).assign(&normed);
        }

        Ok(output)
    }
}

// ---------------------------------------------------------------------------
// Feed-forward network
// ---------------------------------------------------------------------------

/// Feed-forward network sub-layer within a BERT encoder layer.
///
/// Linear -> GELU -> Linear, all with bias terms.
#[derive(Debug)]
struct CpuBertFfn {
    /// Intermediate projection weight `[intermediate, hidden]`.
    intermediate_weight: Array2<f32>,
    /// Intermediate projection bias `[intermediate]`.
    intermediate_bias: Option<Array1<f32>>,
    /// Output projection weight `[hidden, intermediate]`.
    output_weight: Array2<f32>,
    /// Output projection bias `[hidden]`.
    output_bias: Option<Array1<f32>>,
    /// Post-FFN `LayerNorm` weight `[hidden]`.
    output_ln_weight: Array1<f32>,
    /// Post-FFN `LayerNorm` bias `[hidden]`.
    output_ln_bias: Array1<f32>,
    /// Layer norm epsilon.
    layer_norm_eps: f32,
}

impl CpuBertFfn {
    /// FFN forward pass: intermediate projection -> GELU -> output projection -> residual + `LayerNorm`.
    ///
    /// `hidden` is `[seq, hidden]`.
    fn forward(&self, hidden: &Array2<f32>) -> Array2<f32> {
        let seq = hidden.shape()[0];
        let hidden_dim = hidden.shape()[1];

        // Intermediate projection: [seq, hidden] @ [hidden, intermediate] => [seq, intermediate]
        let intermediate = hidden.dot(&self.intermediate_weight.t());
        let intermediate = if let Some(ref bias) = self.intermediate_bias {
            intermediate + bias
        } else {
            intermediate
        };

        let mut activated = intermediate;
        gelu_inplace(&mut activated);

        // Output projection: [seq, intermediate] @ [intermediate, hidden] => [seq, hidden]
        let output = activated.dot(&self.output_weight.t());
        let output = if let Some(ref bias) = self.output_bias {
            output + bias
        } else {
            output
        };

        // Residual + LayerNorm
        let residual = hidden + &output;
        let mut result = Array2::<f32>::zeros((seq, hidden_dim));
        for t in 0..seq {
            let normed = layer_norm(
                &residual.row(t),
                &self.output_ln_weight,
                &self.output_ln_bias,
                self.layer_norm_eps,
            );
            result.row_mut(t).assign(&normed);
        }

        result
    }
}

// ---------------------------------------------------------------------------
// Encoder layer
// ---------------------------------------------------------------------------

/// A single BERT encoder layer (self-attention + FFN).
#[derive(Debug)]
struct CpuBertLayer {
    /// Self-attention sub-layer.
    attention: CpuBertSelfAttention,
    /// Feed-forward sub-layer.
    ffn: CpuBertFfn,
}

impl CpuBertLayer {
    /// Run attention then FFN, returning updated hidden states `[seq, hidden]`.
    fn forward(&self, hidden: &Array2<f32>, mask: &Array1<f32>) -> crate::Result<Array2<f32>> {
        let after_attn = self.attention.forward(hidden, mask)?;
        Ok(self.ffn.forward(&after_attn))
    }
}

// ---------------------------------------------------------------------------
// Full model
// ---------------------------------------------------------------------------

/// Complete BERT model for embedding extraction.
///
/// `pub(super)` so [`CpuRerankBackend`] in this module (and only this
/// module) can hold one. The model has no public API outside the
/// backend types that wrap it.
#[derive(Debug)]
pub(super) struct CpuBertModel {
    /// Embeddings layer (word + position + `token_type` + `LayerNorm`).
    embeddings: CpuBertEmbeddings,
    /// Transformer encoder layers.
    layers: Vec<CpuBertLayer>,
}

impl CpuBertModel {
    /// Run the full BERT forward pass for a single sequence.
    ///
    /// Chains: embeddings -> N attention+FFN layers -> hidden states `[seq, hidden]`.
    fn forward(&self, encoding: &Encoding, mask: &Array1<f32>) -> crate::Result<Array2<f32>> {
        let batched = self.embeddings.forward(std::slice::from_ref(encoding));
        // embeddings.forward returns Vec<Array2> — take the single item
        let mut hidden = batched.into_iter().next().ok_or_else(|| {
            crate::Error::Other(anyhow::anyhow!("embeddings produced empty output"))
        })?;

        for layer in &self.layers {
            hidden = layer.forward(&hidden, mask)?;
        }

        Ok(hidden)
    }
}

/// Load `ClassicBert` encoder layers from safetensors.
fn load_classic_layer(
    tensors: &SafeTensors<'_>,
    i: i32,
    config: &BertConfig,
    name_prefix: &str,
) -> crate::Result<(CpuBertSelfAttention, CpuBertFfn)> {
    let prefix = format!("{name_prefix}encoder.layer.{i}");

    // Load separate Q/K/V weights then fuse via concatenation
    let q_weight = load_tensor2(tensors, &format!("{prefix}.attention.self.query.weight"))?;
    let k_weight = load_tensor2(tensors, &format!("{prefix}.attention.self.key.weight"))?;
    let v_weight = load_tensor2(tensors, &format!("{prefix}.attention.self.value.weight"))?;
    let qkv_weight = ndarray::concatenate(
        Axis(0),
        &[q_weight.view(), k_weight.view(), v_weight.view()],
    )
    .map_err(|e| crate::Error::Other(anyhow::anyhow!("QKV concat error layer {i}: {e}")))?;

    // Fuse biases if present
    let q_bias = try_load_tensor1(tensors, &format!("{prefix}.attention.self.query.bias"))?;
    let k_bias = try_load_tensor1(tensors, &format!("{prefix}.attention.self.key.bias"))?;
    let v_bias = try_load_tensor1(tensors, &format!("{prefix}.attention.self.value.bias"))?;
    let qkv_bias = match (&q_bias, &k_bias, &v_bias) {
        (Some(qb), Some(kb), Some(vb)) => Some(
            ndarray::concatenate(Axis(0), &[qb.view(), kb.view(), vb.view()]).map_err(|e| {
                crate::Error::Other(anyhow::anyhow!("QKV bias concat error layer {i}: {e}"))
            })?,
        ),
        _ => None,
    };

    let attention = CpuBertSelfAttention {
        qkv_weight,
        qkv_bias,
        output_weight: load_tensor2(tensors, &format!("{prefix}.attention.output.dense.weight"))?,
        output_bias: try_load_tensor1(tensors, &format!("{prefix}.attention.output.dense.bias"))?,
        output_ln_weight: load_tensor1(
            tensors,
            &format!("{prefix}.attention.output.LayerNorm.weight"),
        )?,
        output_ln_bias: load_tensor1(
            tensors,
            &format!("{prefix}.attention.output.LayerNorm.bias"),
        )?,
        num_heads: config.num_attention_heads,
        head_dim: config.hidden_size / config.num_attention_heads,
        layer_norm_eps: config.layer_norm_eps,
    };
    let ffn = CpuBertFfn {
        intermediate_weight: load_tensor2(tensors, &format!("{prefix}.intermediate.dense.weight"))?,
        intermediate_bias: try_load_tensor1(tensors, &format!("{prefix}.intermediate.dense.bias"))?,
        output_weight: load_tensor2(tensors, &format!("{prefix}.output.dense.weight"))?,
        output_bias: try_load_tensor1(tensors, &format!("{prefix}.output.dense.bias"))?,
        output_ln_weight: load_tensor1(tensors, &format!("{prefix}.output.LayerNorm.weight"))?,
        output_ln_bias: load_tensor1(tensors, &format!("{prefix}.output.LayerNorm.bias"))?,
        layer_norm_eps: config.layer_norm_eps,
    };
    Ok((attention, ffn))
}

impl CpuBertModel {
    /// Load model weights from a safetensors file.
    ///
    /// Parses all weights into `ndarray` arrays. Fuses separate Q/K/V
    /// weight matrices into a single `[3*hidden, hidden]` matrix for
    /// `ClassicBert` (matching the MLX backend pattern).
    #[expect(
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss,
        clippy::cast_possible_wrap,
        reason = "hidden_size and num_layers are small positive ints from config"
    )]
    fn from_safetensors(tensors: &SafeTensors<'_>, config: &BertConfig) -> crate::Result<Self> {
        // Variant prefix: bare for ClassicBert (BGE), `bert.` for HF
        // transformers exports (MS-MARCO cross-encoders).
        let p = detect_variant(tensors)?;

        let embeddings = CpuBertEmbeddings {
            word_embeddings: load_tensor2(
                tensors,
                &format!("{p}embeddings.word_embeddings.weight"),
            )?,
            position_embeddings: Some(load_tensor2(
                tensors,
                &format!("{p}embeddings.position_embeddings.weight"),
            )?),
            token_type_embeddings: Some(load_tensor2(
                tensors,
                &format!("{p}embeddings.token_type_embeddings.weight"),
            )?),
            layer_norm_weight: load_tensor1(tensors, &format!("{p}embeddings.LayerNorm.weight"))?,
            layer_norm_bias: load_tensor1(tensors, &format!("{p}embeddings.LayerNorm.bias"))?,
            layer_norm_eps: config.layer_norm_eps,
        };

        let emb_dim = embeddings.word_embeddings.shape()[1] as i32;
        if emb_dim != config.hidden_size {
            return Err(crate::Error::Other(anyhow::anyhow!(
                "model hidden_size mismatch: config says {} but word_embeddings has dim {}",
                config.hidden_size,
                emb_dim
            )));
        }

        let mut layers = Vec::with_capacity(config.num_hidden_layers as usize);
        for i in 0..config.num_hidden_layers {
            let (attention, ffn) = load_classic_layer(tensors, i, config, p)?;
            layers.push(CpuBertLayer { attention, ffn });
        }

        Ok(Self { embeddings, layers })
    }
}

// ---------------------------------------------------------------------------
// Public backend wrapper
// ---------------------------------------------------------------------------

/// CPU-based BERT embedding backend using ndarray + system BLAS.
///
/// Uses [`ndarray`] with system BLAS for matrix operations. All computation
/// runs on the CPU, making this backend portable to any platform with a
/// system BLAS library (`OpenBLAS`, MKL, Accelerate).
///
/// Supports the `ClassicBert` family (BGE models), detected automatically
/// from weight names.
///
/// The inner model is wrapped in [`Arc`] so the backend can be cheaply
/// cloned for per-thread use in rayon.
pub struct CpuBackend {
    /// The BERT model (all weights as ndarray arrays).
    model: Arc<CpuBertModel>,
    /// Hidden dimension for output vector size validation.
    hidden_size: i32,
    /// Maximum sequence length supported by the model.
    max_position_embeddings: i32,
}

impl std::fmt::Debug for CpuBackend {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CpuBackend")
            .field("hidden_size", &self.hidden_size)
            .field("max_position_embeddings", &self.max_position_embeddings)
            .finish_non_exhaustive()
    }
}

impl CpuBackend {
    /// Load a `ClassicBert` (BGE) embedding model from `HuggingFace`.
    ///
    /// Downloads `model.safetensors` and `config.json` on first call;
    /// subsequent calls use the `hf-hub` cache. Returns an error if the
    /// model architecture is not recognized.
    ///
    /// # Errors
    ///
    /// Returns an error if the model cannot be downloaded, the config
    /// cannot be parsed, or the weights fail to load.
    pub fn load(model_repo: &str, _device_hint: &DeviceHint) -> crate::Result<Self> {
        // Report BLAS status and recommend optimal library for this CPU
        let blas = super::blas_info::detect_blas();
        let cpu = super::blas_info::detect_cpu_vendor();
        tracing::info!("CPU backend: {} CPU, {} BLAS", cpu, blas);
        if let Some(tip) = super::blas_info::recommend_blas() {
            eprintln!("[ripvec] {tip}");
        }

        let api = Api::new().map_err(|e| crate::Error::Download(e.to_string()))?;
        let repo = api.model(model_repo.to_string());

        let config_path = repo
            .get("config.json")
            .map_err(|e| crate::Error::Download(e.to_string()))?;
        let weights_path = repo
            .get("model.safetensors")
            .map_err(|e| crate::Error::Download(e.to_string()))?;

        let model_bytes = std::fs::read(&weights_path).map_err(|e| crate::Error::Io {
            path: weights_path.display().to_string(),
            source: e,
        })?;

        let tensors = SafeTensors::deserialize(&model_bytes)
            .map_err(|e| crate::Error::Other(anyhow::anyhow!("safetensors parse error: {e}")))?;
        let _ = detect_variant(&tensors)?;

        let config_str = std::fs::read_to_string(&config_path).map_err(|e| crate::Error::Io {
            path: config_path.display().to_string(),
            source: e,
        })?;
        let config_json: serde_json::Value = serde_json::from_str(&config_str)
            .map_err(|e| crate::Error::Other(anyhow::anyhow!("config parse error: {e}")))?;
        let config = BertConfig::from_json(&config_json)?;

        let hidden_size = config.hidden_size;
        let max_position_embeddings = config.max_position_embeddings;
        let model = CpuBertModel::from_safetensors(&tensors, &config)?;

        Ok(Self {
            model: Arc::new(model),
            hidden_size,
            max_position_embeddings,
        })
    }
}

impl EmbedBackend for CpuBackend {
    /// Embed a batch of pre-tokenized inputs using the full BERT forward pass.
    ///
    /// Runs: embeddings -> N attention+FFN layers -> CLS pooling -> L2 normalize.
    /// Each batch item is processed independently (no cross-batch padding needed).
    fn embed_batch(&self, encodings: &[Encoding]) -> crate::Result<Vec<Vec<f32>>> {
        if encodings.is_empty() {
            return Ok(vec![]);
        }

        let mut results = Vec::with_capacity(encodings.len());
        for enc in encodings {
            // Build attention mask: 0.0 for real tokens, -1e9 for padding
            let mask = Array1::from_vec(
                enc.attention_mask
                    .iter()
                    .map(|&m| if m == 1 { 0.0_f32 } else { -1e9_f32 })
                    .collect(),
            );

            // Full forward pass: embeddings -> layers -> hidden [seq, hidden]
            let hidden = self.model.forward(enc, &mask)?;

            // CLS pooling (first token) + L2 normalize
            let cls = hidden.row(0);
            let norm = cls.mapv(|v| v * v).sum().sqrt().max(1e-12);
            let normalized: Vec<f32> = cls.iter().map(|&v| v / norm).collect();
            results.push(normalized);
        }

        Ok(results)
    }

    /// CPU backend supports cheap cloning via `Arc`.
    fn supports_clone(&self) -> bool {
        true
    }

    /// Clone the backend for per-thread use in rayon.
    fn clone_backend(&self) -> Box<dyn EmbedBackend> {
        Box::new(Self {
            model: Arc::clone(&self.model),
            hidden_size: self.hidden_size,
            max_position_embeddings: self.max_position_embeddings,
        })
    }

    /// CPU backend does not use GPU.
    fn is_gpu(&self) -> bool {
        false
    }

    fn name(&self) -> &'static str {
        crate::backend::blas_info::cpu_driver_name()
    }

    /// Maximum tokens from model config (512 for `ClassicBert`).
    #[expect(
        clippy::cast_sign_loss,
        reason = "max_position_embeddings is always positive from config"
    )]
    fn max_tokens(&self) -> usize {
        self.max_position_embeddings as usize
    }
}

// ---------------------------------------------------------------------------
// CPU rerank backend (cross-encoder).
// ---------------------------------------------------------------------------

/// CPU rerank backend: BERT trunk + linear classifier head.
///
/// Shares the trunk with [`CpuBackend`] — the embeddings table, the
/// 12 transformer layers, and the CLS-pooling step are identical to
/// the bi-encoder. Only the tail differs: instead of L2-normalize, the
/// CLS vector is projected through a `Linear(hidden -> 1)` classifier
/// and sigmoid-activated to produce a scalar relevance score.
///
/// Construct via [`Self::load`]. The classifier weights must be
/// present in the model's `model.safetensors`; if not, load returns
/// `Error::Other` ("not a cross-encoder").
pub struct CpuRerankBackend {
    /// Shared BERT trunk. `Arc` for cheap cloning across rayon workers,
    /// matching [`CpuBackend`]'s pattern.
    model: Arc<CpuBertModel>,
    /// BERT pooler weight: `Linear(hidden -> hidden)`, applied to the
    /// CLS row before the classifier head. Stored row-major as
    /// `[hidden_out, hidden_in]` matching the safetensors layout
    /// (PyTorch `nn.Linear` weight convention is `(out, in)`).
    pooler_weight: Array2<f32>,
    /// BERT pooler bias, shape `[hidden]`.
    pooler_bias: Array1<f32>,
    /// Classifier head weights: `Linear(hidden -> 1)`, stored as a
    /// flat `[hidden]` vector since the output dimension is 1.
    classifier_weight: Array1<f32>,
    /// Classifier head bias (scalar).
    classifier_bias: f32,
    /// Max sequence length from `BertConfig::max_position_embeddings`.
    max_position_embeddings: i32,
}

impl std::fmt::Debug for CpuRerankBackend {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CpuRerankBackend")
            .field("max_position_embeddings", &self.max_position_embeddings)
            .field("classifier_dim", &self.classifier_weight.len())
            .finish_non_exhaustive()
    }
}

impl CpuRerankBackend {
    /// Load a cross-encoder model from a `HuggingFace` repo.
    ///
    /// Loads the same `config.json` + `model.safetensors` as
    /// [`CpuBackend::load`], plus the classifier head weights
    /// (`classifier.weight`, `classifier.bias`). Caches via `hf-hub`,
    /// so multiple sub-agent processes share weights through
    /// `~/.cache/huggingface/hub/`.
    ///
    /// # Errors
    ///
    /// Returns an error if downloading fails, the safetensors parse
    /// fails, the model has no classifier head, or the classifier
    /// shape is not `[1, hidden]` (multi-class cross-encoders unsupported).
    pub fn load(model_repo: &str) -> crate::Result<Self> {
        let api = Api::new().map_err(|e| crate::Error::Download(e.to_string()))?;
        let repo = api.model(model_repo.to_string());

        let config_path = repo
            .get("config.json")
            .map_err(|e| crate::Error::Download(e.to_string()))?;
        let weights_path = repo
            .get("model.safetensors")
            .map_err(|e| crate::Error::Download(e.to_string()))?;

        let model_bytes = std::fs::read(&weights_path).map_err(|e| crate::Error::Io {
            path: weights_path.display().to_string(),
            source: e,
        })?;
        let tensors = SafeTensors::deserialize(&model_bytes)
            .map_err(|e| crate::Error::Other(anyhow::anyhow!("safetensors parse error: {e}")))?;
        let _ = detect_variant(&tensors)?;

        let config_str = std::fs::read_to_string(&config_path).map_err(|e| crate::Error::Io {
            path: config_path.display().to_string(),
            source: e,
        })?;
        let config_json: serde_json::Value = serde_json::from_str(&config_str)
            .map_err(|e| crate::Error::Other(anyhow::anyhow!("config parse error: {e}")))?;
        let config = BertConfig::from_json(&config_json)?;
        let max_position_embeddings = config.max_position_embeddings;

        let model = CpuBertModel::from_safetensors(&tensors, &config)?;
        let (pooler_weight, pooler_bias) = load_pooler_head(&tensors, &config)?;
        let (classifier_weight, classifier_bias) = load_classifier_head(&tensors, &config)?;

        Ok(Self {
            model: Arc::new(model),
            pooler_weight,
            pooler_bias,
            classifier_weight,
            classifier_bias,
            max_position_embeddings,
        })
    }
}

impl RerankBackend for CpuRerankBackend {
    /// Score a batch: forward through trunk, take CLS, project
    /// through classifier, sigmoid.
    ///
    /// Pairs run in parallel via rayon (the trunk is `Arc`-shared and
    /// `forward` is pure). Empty input returns empty output.
    fn score_batch(&self, encodings: &[Encoding]) -> crate::Result<Vec<f32>> {
        if encodings.is_empty() {
            return Ok(Vec::new());
        }

        encodings
            .par_iter()
            .map(|enc| {
                // Same attention mask convention as CpuBackend::embed_batch:
                // 0.0 for real tokens, -1e9 (effectively -inf in softmax)
                // for padding.
                let mask = Array1::from_vec(
                    enc.attention_mask
                        .iter()
                        .map(|&m| if m == 1 { 0.0_f32 } else { -1e9_f32 })
                        .collect(),
                );

                // Full trunk forward: embeddings → 12 layers → [seq, hidden]
                let hidden = self.model.forward(enc, &mask)?;

                // CLS pooling — first token's hidden state.
                let cls = hidden.row(0).to_owned();

                // BERT pooler: tanh(W_pool @ cls + b_pool). The classifier
                // weights were trained against this pooled CLS, not the
                // raw trunk output. `pooler_weight.dot(&cls)` computes
                // `W @ cls` (matrix [hidden,hidden] × vector [hidden]),
                // equivalent to PyTorch's `cls @ W.T + b`.
                let pooled_pre = self.pooler_weight.dot(&cls);
                let pooled = (&pooled_pre + &self.pooler_bias).mapv(f32::tanh);

                // classifier(pooled) = pooled · W + b. Returns the raw
                // logit; sentence-transformers' ms-marco config declares
                // `Identity` as the default activation, so the canonical
                // public score for this model is the unsigmoided logit
                // (~[-11, +5] empirically). Normalization (if needed)
                // happens at the call site that consumes the score.
                let logit: f32 = pooled
                    .iter()
                    .zip(self.classifier_weight.iter())
                    .map(|(a, b)| a * b)
                    .sum::<f32>()
                    + self.classifier_bias;
                Ok(logit)
            })
            .collect()
    }

    fn is_gpu(&self) -> bool {
        false
    }

    #[expect(
        clippy::cast_sign_loss,
        reason = "max_position_embeddings is always positive from config"
    )]
    fn max_tokens(&self) -> usize {
        self.max_position_embeddings as usize
    }

    fn name(&self) -> &'static str {
        crate::backend::blas_info::cpu_driver_name()
    }
}

/// Load the BERT pooler (`bert.pooler.dense.{weight,bias}`).
///
/// BERT cross-encoders apply a `tanh(Linear(hidden -> hidden))` to the
/// CLS row before the classification head — the head was trained
/// against this pooled representation, not raw CLS. Skipping it (which
/// ripvec did from day one) leaves the classifier multiplying its
/// trained weights against unpooled CLS activations, producing logits
/// that wander near 0 and squash to ~0.5 under sigmoid.
///
/// Accepts both `bert.pooler.dense.*` (canonical BERT naming) and
/// `pooler.dense.*` (some HF re-exports drop the `bert.` prefix).
fn load_pooler_head(
    tensors: &SafeTensors<'_>,
    config: &BertConfig,
) -> crate::Result<(Array2<f32>, Array1<f32>)> {
    let weight_tensor = tensors
        .tensor("bert.pooler.dense.weight")
        .or_else(|_| tensors.tensor("pooler.dense.weight"))
        .map_err(|_| {
            crate::Error::Other(anyhow::anyhow!(
                "BERT pooler weight not found (looked for \
                 bert.pooler.dense.weight, pooler.dense.weight). \
                 Cross-encoder classification heads need the pooler; \
                 if the model legitimately has no pooler, this code \
                 path will need a feature gate."
            ))
        })?;
    #[expect(
        clippy::cast_sign_loss,
        reason = "hidden_size is a small positive int from config"
    )]
    let hidden = config.hidden_size as usize;
    let shape = weight_tensor.shape();
    if shape.len() != 2 || shape[0] != hidden || shape[1] != hidden {
        return Err(crate::Error::Other(anyhow::anyhow!(
            "pooler weight has shape {shape:?}; expected [{hidden}, {hidden}]"
        )));
    }
    let weight_vec = decode_f32_tensor(weight_tensor.data(), weight_tensor.dtype())?;
    let weight = Array2::from_shape_vec((hidden, hidden), weight_vec)
        .map_err(|e| crate::Error::Other(anyhow::anyhow!("pooler weight reshape: {e}")))?;

    let bias_tensor = tensors
        .tensor("bert.pooler.dense.bias")
        .or_else(|_| tensors.tensor("pooler.dense.bias"))
        .map_err(|_| {
            crate::Error::Other(anyhow::anyhow!(
                "pooler bias not found alongside pooler weight"
            ))
        })?;
    let bias_vec = decode_f32_tensor(bias_tensor.data(), bias_tensor.dtype())?;
    if bias_vec.len() != hidden {
        return Err(crate::Error::Other(anyhow::anyhow!(
            "pooler bias length {} does not match hidden_size {hidden}",
            bias_vec.len()
        )));
    }
    let bias = Array1::from_vec(bias_vec);

    Ok((weight, bias))
}

/// Load the cross-encoder classifier head from the same
/// `model.safetensors` file as the trunk.
///
/// Returns `(weight, bias)` where `weight` is a flat `[hidden]`
/// vector (the `[1, hidden]` 2D tensor's only row) and `bias` is the
/// scalar from the `[1]` bias tensor.
///
/// Tries several common naming conventions: ms-marco checkpoints use
/// `classifier.{weight,bias}`; some HF exports use
/// `model.classifier.*`. Errors clearly when no head is found,
/// distinguishing "this is a bi-encoder by mistake" from generic
/// safetensors errors.
fn load_classifier_head(
    tensors: &SafeTensors<'_>,
    config: &BertConfig,
) -> crate::Result<(Array1<f32>, f32)> {
    let weight_tensor = tensors
        .tensor("classifier.weight")
        .or_else(|_| tensors.tensor("model.classifier.weight"))
        .map_err(|_| {
            crate::Error::Other(anyhow::anyhow!(
                "no classifier head found in safetensors — is this a \
                 cross-encoder model, or a bi-encoder by mistake?"
            ))
        })?;
    let shape = weight_tensor.shape();
    if shape.len() != 2 || shape[0] != 1 {
        return Err(crate::Error::Other(anyhow::anyhow!(
            "classifier weight has shape {shape:?}; expected [1, hidden]. \
             Multi-class cross-encoders are not supported (ripvec uses \
             single-score relevance)."
        )));
    }
    #[expect(
        clippy::cast_sign_loss,
        reason = "hidden_size is a small positive int from config"
    )]
    let expected_hidden = config.hidden_size as usize;
    if shape[1] != expected_hidden {
        return Err(crate::Error::Other(anyhow::anyhow!(
            "classifier hidden dim {} does not match config hidden_size {expected_hidden}",
            shape[1]
        )));
    }

    let weight = decode_f32_tensor(weight_tensor.data(), weight_tensor.dtype())?;
    let weight = Array1::from_vec(weight);

    let bias_tensor = tensors
        .tensor("classifier.bias")
        .or_else(|_| tensors.tensor("model.classifier.bias"))
        .map_err(|_| {
            crate::Error::Other(anyhow::anyhow!(
                "classifier.bias not found alongside classifier.weight"
            ))
        })?;
    let bias_vec = decode_f32_tensor(bias_tensor.data(), bias_tensor.dtype())?;
    let bias = *bias_vec
        .first()
        .ok_or_else(|| crate::Error::Other(anyhow::anyhow!("classifier bias is empty")))?;

    Ok((weight, bias))
}

/// Decode a flat tensor blob into `Vec<f32>`, accepting F32 / F16 /
/// BF16. Used by the classifier-head loader; the trunk uses dedicated
/// `load_tensor*` helpers above.
fn decode_f32_tensor(raw: &[u8], dtype: safetensors::tensor::Dtype) -> crate::Result<Vec<f32>> {
    use safetensors::tensor::Dtype;
    match dtype {
        Dtype::F32 => Ok(raw
            .chunks_exact(4)
            .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
            .collect()),
        Dtype::F16 => Ok(raw
            .chunks_exact(2)
            .map(|b| half::f16::from_le_bytes([b[0], b[1]]).to_f32())
            .collect()),
        Dtype::BF16 => Ok(raw
            .chunks_exact(2)
            .map(|b| half::bf16::from_le_bytes([b[0], b[1]]).to_f32())
            .collect()),
        other => Err(crate::Error::Other(anyhow::anyhow!(
            "unsupported classifier dtype: {other:?}"
        ))),
    }
}

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

    const BGE_SMALL: &str = "BAAI/bge-small-en-v1.5";

    #[test]
    fn config_from_json_classic() {
        let json: serde_json::Value = serde_json::json!({
            "hidden_size": 384,
            "num_hidden_layers": 12,
            "num_attention_heads": 12,
            "max_position_embeddings": 512,
            "layer_norm_eps": 1e-12
        });
        let config = BertConfig::from_json(&json).unwrap();
        assert_eq!(config.hidden_size, 384);
        assert_eq!(config.num_hidden_layers, 12);
        assert_eq!(config.num_attention_heads, 12);
        assert_eq!(config.max_position_embeddings, 512);
    }

    #[test]
    fn config_missing_key_errors() {
        let json: serde_json::Value = serde_json::json!({});
        let result = BertConfig::from_json(&json);
        assert!(result.is_err());
    }

    #[test]
    fn cpu_backend_loads_model() {
        let backend = CpuBackend::load(BGE_SMALL, &DeviceHint::Cpu).unwrap();
        assert_eq!(backend.hidden_size, 384);
        assert_eq!(backend.max_position_embeddings, 512);
        assert!(!backend.is_gpu());
        assert!(backend.supports_clone());
        assert_eq!(backend.max_tokens(), 512);
    }

    #[test]
    fn cpu_backend_embeddings_forward() {
        let backend = CpuBackend::load(BGE_SMALL, &DeviceHint::Cpu).unwrap();
        let enc = Encoding {
            input_ids: vec![101, 2023, 2003, 1037, 3231, 102],
            attention_mask: vec![1, 1, 1, 1, 1, 1],
            token_type_ids: vec![0, 0, 0, 0, 0, 0],
        };
        let outputs = backend.model.embeddings.forward(&[enc]);
        assert_eq!(outputs.len(), 1);
        assert_eq!(outputs[0].shape(), &[6, 384]);
        let sum: f32 = outputs[0].iter().map(|v| v.abs()).sum();
        assert!(sum > 0.0, "embeddings output should not be all zeros");
    }

    #[test]
    fn cpu_backend_clone() {
        let backend = CpuBackend::load(BGE_SMALL, &DeviceHint::Cpu).unwrap();
        let cloned = backend.clone_backend();
        assert!(!cloned.is_gpu());
        assert!(cloned.supports_clone());
        assert_eq!(cloned.max_tokens(), 512);
    }

    #[test]
    fn cpu_backend_full_forward_output_dim() {
        // "hello world" tokenized: [CLS]=101, hello=7592, world=2088, [SEP]=102
        let backend = CpuBackend::load(BGE_SMALL, &DeviceHint::Cpu).unwrap();
        let enc = Encoding {
            input_ids: vec![101, 7592, 2088, 102],
            attention_mask: vec![1, 1, 1, 1],
            token_type_ids: vec![0, 0, 0, 0],
        };
        let result = backend.embed_batch(&[enc]).unwrap();
        assert_eq!(result.len(), 1);
        assert_eq!(
            result[0].len(),
            384,
            "BGE-small should produce 384-dim embeddings"
        );
    }

    #[test]
    fn cpu_backend_full_forward_l2_norm() {
        let backend = CpuBackend::load(BGE_SMALL, &DeviceHint::Cpu).unwrap();
        let enc = Encoding {
            input_ids: vec![101, 7592, 2088, 102],
            attention_mask: vec![1, 1, 1, 1],
            token_type_ids: vec![0, 0, 0, 0],
        };
        let result = backend.embed_batch(&[enc]).unwrap();
        let norm: f32 = result[0].iter().map(|v| v * v).sum::<f32>().sqrt();
        assert!(
            (norm - 1.0).abs() < 1e-4,
            "L2 norm should be ~1.0, got {norm}"
        );
    }

    #[test]
    fn cpu_backend_different_inputs_differ() {
        let backend = CpuBackend::load(BGE_SMALL, &DeviceHint::Cpu).unwrap();
        let enc1 = Encoding {
            input_ids: vec![101, 7592, 2088, 102], // "hello world"
            attention_mask: vec![1, 1, 1, 1],
            token_type_ids: vec![0, 0, 0, 0],
        };
        let enc2 = Encoding {
            input_ids: vec![101, 19387, 8840, 4313, 102], // "quantum physics"
            attention_mask: vec![1, 1, 1, 1, 1],
            token_type_ids: vec![0, 0, 0, 0, 0],
        };
        let results = backend.embed_batch(&[enc1, enc2]).unwrap();
        assert_eq!(results.len(), 2);

        // Cosine similarity (both L2-normalized, so dot product = cosine sim)
        let dot: f32 = results[0]
            .iter()
            .zip(results[1].iter())
            .map(|(a, b)| a * b)
            .sum();
        assert!(
            dot < 0.99,
            "different inputs should produce different embeddings, cosine sim = {dot}"
        );
    }

    #[test]
    fn cpu_backend_empty_batch() {
        let backend = CpuBackend::load(BGE_SMALL, &DeviceHint::Cpu).unwrap();
        let result = backend.embed_batch(&[]).unwrap();
        assert!(result.is_empty(), "empty batch should return empty vec");
    }
}
