candding 0.1.0

use std::path::PathBuf;

use candle_core::Tensor;

use super::*;
use crate::models::EncoderInput;
use crate::pooling::Pooling;
use crate::registry;
use crate::registry::BackendSupport;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::tokenizer::PadToken;

fn assert_send_sync<T: Send + Sync>() {}

/// The golden suite caches one model per test module in a `OnceLock`, which needs `Send + Sync`.
#[test]
fn text_embedding_is_send_and_sync() {
  assert_send_sync::<TextEmbedding>();
}

/// Counting reads the tokenizer alone, so the encoder never runs.
struct StubEncoder;

impl Encoder for StubEncoder {
  fn forward(&self, _input: &EncoderInput) -> Result<Tensor> {
    unreachable!("counting never runs the encoder")
  }

  fn hidden_size(&self) -> usize {
    1024
  }

  fn max_position_embeddings(&self) -> usize {
    8192
  }

  fn weights_dtype(&self) -> DType {
    DType::F32
  }
}

fn counting_model(id: &str) -> TextEmbedding {
  let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/tiny_tokenizer.json");
  TextEmbedding::new(
    Box::new(StubEncoder),
    TextTokenizer::from_file(&path, 512, PadToken::Id(0)).unwrap(),
    Vec::new(),
    registry::dense::lookup(id).unwrap(),
    Device::Cpu,
    DType::F32,
    32,
  )
}

#[test]
fn a_custom_instruction_changes_the_query_token_count() {
  let qwen = counting_model("Qwen/Qwen3-Embedding-0.6B");
  let short = "Find x.";
  let long =
    "Given a web search query, retrieve relevant passages that answer the query in full detail.";
  let with_short = qwen
    .query_token_count_with_instruction(&["hello world"], short)
    .unwrap();
  let with_long = qwen
    .query_token_count_with_instruction(&["hello world"], long)
    .unwrap();
  assert!(
    with_long[0] > with_short[0],
    "{with_long:?} against {with_short:?}"
  );
  let templated = template::apply(
    qwen.descriptor().query_template.as_deref(),
    Some(short),
    "hello world",
  );
  assert_eq!(
    with_short,
    qwen.token_count(&[templated]).unwrap(),
    "the count is the count of the text query_embed_with_instruction encodes"
  );
}

#[test]
fn a_template_changes_the_token_count() {
  let qwen = counting_model("Qwen/Qwen3-Embedding-0.6B");
  let plain = qwen.token_count(&["hello world"]).unwrap();
  assert_eq!(plain, vec![4], "[CLS] hello world [SEP]");
  let query = qwen.query_token_count(&["hello world"]).unwrap();
  assert!(
    query[0] > plain[0],
    "the instruct template adds tokens: {query:?} against {plain:?}"
  );
  let templated = template::apply(
    qwen.descriptor().query_template.as_deref(),
    qwen.descriptor().default_instruction.as_deref(),
    "hello world",
  );
  assert_eq!(
    query,
    qwen.token_count(&[templated]).unwrap(),
    "the count is the count of the text query_embed encodes"
  );
  assert_eq!(
    qwen.passage_token_count(&["hello world"]).unwrap(),
    plain,
    "this entry has no passage template"
  );

  let e5 = counting_model("intfloat/multilingual-e5-small");
  assert_eq!(e5.token_count(&["hello world"]).unwrap(), vec![4]);
  for counts in [
    e5.query_token_count(&["hello world"]).unwrap(),
    e5.passage_token_count(&["hello world"]).unwrap(),
  ] {
    assert_eq!(
      counts,
      vec![6],
      "the prefix and its colon are two more tokens on either side"
    );
  }
}

/// Returns 3.0 everywhere, so a pooled row's raw L2 norm is `(4.0 * 9.0f32).sqrt()`.
struct ConstantEncoder;

impl Encoder for ConstantEncoder {
  fn forward(&self, input: &EncoderInput) -> Result<Tensor> {
    let (batch, seq) = input.input_ids.dims2()?;
    Ok(Tensor::full(
      3.0f32,
      (batch, seq, self.hidden_size()),
      input.input_ids.device(),
    )?)
  }

  fn hidden_size(&self) -> usize {
    4
  }

  fn max_position_embeddings(&self) -> usize {
    512
  }

  fn weights_dtype(&self) -> DType {
    DType::F32
  }
}

/// `override_` stands in for the builder's `normalize` field: `None` keeps `descriptor_normalize`, `Some` replaces it.
fn constant_model(descriptor_normalize: bool, override_: Option<bool>) -> TextEmbedding {
  let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/tiny_tokenizer.json");
  let descriptor = ModelDescriptor {
    id: "stub".to_string(),
    family: Family::Bert,
    dim: 4,
    default_max_length: 512,
    pooling: Pooling::Cls,
    normalize: override_.unwrap_or(descriptor_normalize),
    query_template: None,
    passage_template: None,
    default_instruction: None,
    dense_dirs: Vec::new(),
    mrl: false,
    layer_norm_after_pooling: false,
    license: "unknown".to_string(),
    size_gb: 0.0,
    gated: false,
    phase: u8::MAX,
    backends: BackendSupport::untested(),
    default_dtype: DType::F32,
    dtype_backends: DtypeBackendSupport::untested(),
  };
  TextEmbedding::new(
    Box::new(ConstantEncoder),
    TextTokenizer::from_file(&path, 512, PadToken::Id(0)).unwrap(),
    Vec::new(),
    descriptor,
    Device::Cpu,
    DType::F32,
    32,
  )
}

fn l2_norm(row: &[f32]) -> f32 {
  row.iter().map(|v| v * v).sum::<f32>().sqrt()
}

#[test]
fn the_descriptors_normalize_flag_controls_whether_rows_are_unit_length() {
  let raw_norm = (4.0 * 9.0f32).sqrt();

  let unit = constant_model(true, None)
    .embed(&["hello world"], None)
    .unwrap();
  assert!((l2_norm(&unit[0]) - 1.0).abs() < 1e-6, "{:?}", unit[0]);

  let raw = constant_model(true, Some(false))
    .embed(&["hello world"], None)
    .unwrap();
  assert!((l2_norm(&raw[0]) - raw_norm).abs() < 1e-6, "{:?}", raw[0]);

  let overridden = constant_model(false, Some(true))
    .embed(&["hello world"], None)
    .unwrap();
  assert!(
    (l2_norm(&overridden[0]) - 1.0).abs() < 1e-6,
    "{:?}",
    overridden[0]
  );
}