candding 0.1.0

use std::path::PathBuf;

use super::*;
use crate::error::CanddingError;
use crate::models::MultiVectorHead;
use crate::registry::BackendSupport;
use crate::registry::Catalog;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::registry::Support;
use crate::registry::WeightsFormat;
use crate::tokenizer::PadToken;

mod guards;
mod markers;

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 multi_vector_text_embedding_is_send_and_sync() {
  assert_send_sync::<MultiVectorTextEmbedding>();
}

/// Reads no tensors: each embedding is one row, `[real-token length, rows in this forward call]`.
struct StubMultiVectorModel;

impl MultiVectorModel for StubMultiVectorModel {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>> {
    let chunk_len = input.lengths.len() as f32;
    Ok(
      input
        .lengths
        .iter()
        .map(|&len| MultiVectorEmbedding::from_rows(2, [vec![len as f32, chunk_len]]))
        .collect(),
    )
  }

  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>> {
    self.forward(input)
  }

  fn dim(&self) -> usize {
    2
  }

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

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

fn stub_descriptor() -> MultiVectorDescriptor {
  MultiVectorDescriptor {
    id: "stub/multi-vector".to_string(),
    family: Family::Bert,
    head: MultiVectorHead::M3Colbert,
    default_max_length: 512,
    keeps_cls_row: false,
    projection_in: 2,
    projection_out: 2,
    projection_bias: true,
    normalize: false,
    score_reduction: crate::multi_vector::ScoreReduction::Sum,
    query_marker: None,
    document_marker: None,
    query_length: None,
    attends_to_mask_tokens: false,
    masked_document_row_fill: crate::multi_vector::MaskedRowFill::DtypeMinimum,
    position_scheme: None,
    backends: BackendSupport::untested(),
    dtype_backends: DtypeBackendSupport::untested(),
    license: "unknown".to_string(),
    size_gb: 0.0,
    phase: u8::MAX,
    weights: WeightsFormat::Safetensors,
    gated: false,
  }
}

fn stub_model(max_length: usize, batch_size: usize) -> MultiVectorTextEmbedding {
  let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/tiny_tokenizer.json");
  MultiVectorTextEmbedding::new(
    Box::new(StubMultiVectorModel),
    TextTokenizer::from_file(&path, max_length, PadToken::Id(0)).unwrap(),
    TextTokenizer::from_file(&path, max_length, PadToken::Id(0)).unwrap(),
    stub_descriptor(),
    Device::Cpu,
    DType::F32,
    batch_size,
  )
}

/// "a", "a b" and "a b c" tokenize to 3, 4 and 5 tokens, so batch size 2 sends two rows, then one.
#[test]
fn batch_size_chunks_the_input_and_results_come_back_in_order() {
  let model = stub_model(512, 2);
  let texts = ["a", "a b c", "a b"];
  let rows = model.embed(&texts, None).unwrap();
  assert_eq!(rows.len(), 3);
  let row = |i: usize| rows[i].rows().next().unwrap().to_vec();
  assert_eq!(
    row(0),
    vec![3.0, 1.0],
    "text 0, \"a\", was the lone second chunk"
  );
  assert_eq!(
    row(1),
    vec![5.0, 2.0],
    "text 1, \"a b c\", shared the first chunk"
  );
  assert_eq!(
    row(2),
    vec![4.0, 2.0],
    "text 2, \"a b\", shared the first chunk"
  );
}

#[test]
fn max_length_reaches_the_tokenizer() {
  assert_eq!(stub_model(16, 8).max_length(), 16);
  assert_eq!(stub_model(64, 8).max_length(), 64);
}

/// `tiny_tokenizer_no_special_tokens.json` gives the empty string zero tokens.
fn stub_model_no_special_tokens(max_length: usize, batch_size: usize) -> MultiVectorTextEmbedding {
  let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
    .join("tests/fixtures/tiny_tokenizer_no_special_tokens.json");
  MultiVectorTextEmbedding::new(
    Box::new(StubMultiVectorModel),
    TextTokenizer::from_file(&path, max_length, PadToken::Id(0)).unwrap(),
    TextTokenizer::from_file(&path, max_length, PadToken::Id(0)).unwrap(),
    stub_descriptor(),
    Device::Cpu,
    DType::F32,
    batch_size,
  )
}

/// The stub returns a row even for a zero-length text, so an empty embedding proves the model never ran.
#[test]
fn an_all_empty_chunk_never_reaches_the_model() {
  let model = stub_model_no_special_tokens(512, 8);
  let rows = model.embed(&[""], None).unwrap();
  assert_eq!(rows.len(), 1);
  assert!(
    rows[0].is_empty(),
    "an all-empty chunk must short-circuit before the model runs: {:?}",
    rows[0]
  );
}

#[test]
fn a_mixed_chunk_still_reaches_the_model() {
  let model = stub_model_no_special_tokens(512, 8);
  let rows = model.embed(&["a", ""], None).unwrap();
  assert_eq!(rows.len(), 2);
  let row = |i: usize| rows[i].rows().next().unwrap().to_vec();
  assert_eq!(
    row(0),
    vec![1.0, 2.0],
    "\"a\" is one token with no special tokens added"
  );
  assert_eq!(
    row(1),
    vec![0.0, 2.0],
    "the empty row still reaches the model"
  );
}

#[test]
fn an_unregistered_id_is_refused_by_name() {
  match MultiVectorTextEmbedding::builder("org/definitely-not-registered").build() {
    Err(CanddingError::UnregisteredModel { id, catalog }) => {
      assert_eq!(id, "org/definitely-not-registered");
      assert_eq!(catalog, Catalog::MultiVector);
    }
    Ok(_) => panic!("an unregistered id must be refused"),
    Err(other) => panic!("unexpected error: {other}"),
  }
}