candding 0.1.0

use std::path::PathBuf;

use super::*;
use crate::error::CanddingError;
use crate::error::LengthKind;
use crate::models::RerankHead;
use crate::registry::BackendSupport;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::registry::WeightsFormat;
use crate::registry::rerank::RerankInput;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PadToken;
use crate::tokenizer::PairTruncationPolicy;

mod prompt;
mod refusals;
mod windowed;

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_cross_encoder_is_send_and_sync() {
  assert_send_sync::<TextCrossEncoder>();
}

/// Reads no tensors and scores each row by its real-token length.
struct StubRerankModel;

impl RerankModel for StubRerankModel {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<RerankScore>> {
    Ok(
      input
        .lengths
        .iter()
        .map(|&len| RerankScore::new(len as f32))
        .collect(),
    )
  }

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

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

struct PanicsIfCalled;

impl RerankModel for PanicsIfCalled {
  fn forward(&self, _input: &EncoderInput) -> Result<Vec<RerankScore>> {
    panic!("RerankModel::forward must not be called")
  }

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

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

fn stub_descriptor() -> RerankDescriptor {
  RerankDescriptor {
    id: "stub/rerank".to_string(),
    family: Family::Bert,
    head: RerankHead::BertPooler,
    default_max_length: 512,
    score_activation: ScoreActivation::Identity,
    input: RerankInput::Pair {
      truncation: PairTruncationPolicy::LongestFirst,
    },
    alibi_variant: None,
    position_scheme: None,
    default_dtype: DType::F32,
    backends: BackendSupport::untested(),
    dtype_backends: DtypeBackendSupport::untested(),
    license: "unknown".to_string(),
    size_gb: 0.0,
    phase: u8::MAX,
    weights: WeightsFormat::Safetensors,
    gated: false,
  }
}

fn tiny_tokenizer_path() -> PathBuf {
  PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/tiny_tokenizer.json")
}

fn stub_model(
  model: Box<dyn RerankModel>,
  max_length: usize,
  batch_size: usize,
) -> TextCrossEncoder {
  stub_model_with(
    model,
    max_length,
    batch_size,
    PairTruncationPolicy::LongestFirst,
  )
}

fn stub_model_with(
  model: Box<dyn RerankModel>,
  max_length: usize,
  batch_size: usize,
  pair_truncation: PairTruncationPolicy,
) -> TextCrossEncoder {
  let mut descriptor = stub_descriptor();
  descriptor.input = RerankInput::Pair {
    truncation: pair_truncation,
  };
  cross_encoder(
    model,
    descriptor,
    &tiny_tokenizer_path(),
    max_length,
    batch_size,
  )
}

fn cross_encoder(
  model: Box<dyn RerankModel>,
  descriptor: RerankDescriptor,
  tokenizer: &std::path::Path,
  max_length: usize,
  batch_size: usize,
) -> TextCrossEncoder {
  let tokenizer =
    InputTokenizer::from_file(&descriptor.input, tokenizer, max_length, PadToken::Id(0)).unwrap();
  TextCrossEncoder::new(
    model,
    tokenizer,
    descriptor,
    Device::Cpu,
    DType::F32,
    batch_size,
  )
}

/// ("a", "a"), ("a", "a b") and ("a", "a b c") are 5, 6 and 7 tokens as `[CLS] q [SEP] d [SEP]`.
#[test]
fn rerank_builds_a_real_pair_per_document_and_keeps_the_callers_order() {
  let model = stub_model(Box::new(StubRerankModel), 512, 8);
  let documents = ["a", "a b c", "a b"];
  let scores = model.rerank("a", &documents, None).unwrap();
  assert_eq!(scores.len(), 3);
  assert_eq!(scores[0].value(), 5.0, "(\"a\", \"a\")");
  assert_eq!(scores[1].value(), 7.0, "(\"a\", \"a b c\")");
  assert_eq!(scores[2].value(), 6.0, "(\"a\", \"a b\")");
}

#[test]
fn batch_size_chunks_the_document_list() {
  let documents = ["a", "a b", "a b c", "a"];
  for batch_size in [1, 2, 3, 8] {
    let model = stub_model(Box::new(StubRerankModel), 512, batch_size);
    let scores = model.rerank("a", &documents, None).unwrap();
    let values: Vec<f32> = scores.iter().map(RerankScore::value).collect();
    assert_eq!(values, vec![5.0, 6.0, 7.0, 5.0], "batch_size {batch_size}");
  }
}

#[test]
fn an_empty_document_list_never_reaches_the_model() {
  let model = stub_model(Box::new(PanicsIfCalled), 512, 8);
  let documents: [&str; 0] = [];
  let scores = model.rerank("a query", &documents, None).unwrap();
  assert!(scores.is_empty());
}

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

#[test]
fn pair_token_count_matches_what_rerank_hands_the_model_under_both_policies() {
  let long = vec!["b"; 30].join(" ");
  let documents = ["a", "a b c", "", long.as_str()];
  for query in ["a b".to_string(), vec!["a"; 20].join(" ")] {
    for policy in [
      PairTruncationPolicy::LongestFirst,
      PairTruncationPolicy::PreserveQueryToThreeQuarters,
    ] {
      let model = stub_model_with(Box::new(StubRerankModel), 20, 8, policy);
      let counts = model.pair_token_count(&query, &documents).unwrap();
      let scored: Vec<usize> = model
        .rerank(&query, &documents, None)
        .unwrap()
        .iter()
        .map(|score| score.value() as usize)
        .collect();
      assert_eq!(counts, scored, "{policy:?}, query {query:?}");
    }
  }
  let model = stub_model(Box::new(PanicsIfCalled), 512, 8);
  let counts = model.pair_token_count("a", &["a", "a b c"]).unwrap();
  assert_eq!(counts, vec![5, 7]);
}

/// At a budget of 20, `LongestFirst` keeps 16 + 1 + 3 tokens and `PreserveQueryToThreeQuarters` 15 + 1 + 3.
#[test]
fn the_descriptors_pair_truncation_reaches_both_the_count_and_the_scored_batch() {
  let query = vec!["a"; 20].join(" ");
  for (policy, expected) in [
    (PairTruncationPolicy::PreserveQueryToThreeQuarters, 19),
    (PairTruncationPolicy::LongestFirst, 20),
  ] {
    let model = stub_model_with(Box::new(StubRerankModel), 20, 8, policy);
    assert_eq!(
      model.pair_token_count(&query, &["a"]).unwrap(),
      vec![expected],
      "{policy:?} pair_token_count"
    );
    let scores = model.rerank(&query, &["a"], None).unwrap();
    assert_eq!(scores[0].value(), expected as f32, "{policy:?} rerank");
  }
}

/// Two special tokens for a text and three for a pair, so a budget of 3 fits a text and no pair.
#[test]
fn a_max_length_with_no_room_for_a_pair_is_refused_when_the_tokenizer_is_built() {
  let built = InputTokenizer::from_file(
    &stub_descriptor().input,
    &tiny_tokenizer_path(),
    3,
    PadToken::Id(0),
  );
  match built {
    Err(
      err @ CanddingError::LengthTooSmall {
        minimum: 4,
        kind: LengthKind::Pair { added: 3 },
        ..
      },
    ) => {
      assert!(err.to_string().contains("pair"), "{err}");
    }
    Err(other) => panic!("unexpected error: {other}"),
    Ok(_) => panic!("max_length 3 leaves a pair no room"),
  }
}