candding 0.1.0

use std::path::PathBuf;

use super::*;
use crate::error::CanddingError;
use crate::models::SparseHead;
use crate::registry::BackendSupport;
use crate::registry::Catalog;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::registry::WeightsFormat;
use crate::registry::sparse::SparseComputation;
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 sparse_text_embedding_is_send_and_sync() {
  assert_send_sync::<SparseTextEmbedding>();
}

/// Reads no tensors: each row is `(real-token length, rows in this forward call)`.
struct StubSparseModel;

impl SparseModel for StubSparseModel {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let chunk_len = input.lengths.len() as f32;
    Ok(
      input
        .lengths
        .iter()
        .map(|&len| SparseEmbedding::from_pairs([(len as u32, chunk_len)]))
        .collect(),
    )
  }

  /// The passage rows with every weight negated, so a test can tell which side a call reached.
  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    let rows = self.forward(input)?;
    let negate = |row: SparseEmbedding| {
      SparseEmbedding::from_pairs(
        row
          .indices
          .into_iter()
          .zip(row.values.into_iter().map(|v| -v)),
      )
    };
    Ok(rows.into_iter().map(negate).collect())
  }

  fn vocab_size(&self) -> usize {
    16
  }

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

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

fn stub_descriptor() -> SparseDescriptor {
  SparseDescriptor {
    id: "stub/sparse".to_string(),
    family: Some(Family::Bert),
    head: SparseHead::SpladeMlm,
    vocab_size: 16,
    default_max_length: 512,
    backends: BackendSupport::untested(),
    dtype_backends: DtypeBackendSupport::untested(),
    license: "unknown".to_string(),
    size_gb: 0.0,
    phase: u8::MAX,
    weights: WeightsFormat::Safetensors,
    gated: false,
    head_source: None,
    head_revision: None,
    computation: SparseComputation::Encoded,
  }
}

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

#[test]
fn only_query_embed_reaches_the_query_side() {
  let model = stub_model(512, 8);
  let passage = model.passage_embed(&["a b"], None).unwrap();
  assert_eq!(model.embed(&["a b"], None).unwrap(), passage);
  assert_eq!(passage[0].values, [1.0]);
  assert_eq!(model.query_embed(&["a b"], None).unwrap()[0].values, [-1.0]);
}

/// "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);
  assert_eq!(rows[0].indices, vec![3], "text 0, \"a\", is 3 tokens");
  assert_eq!(rows[1].indices, vec![5], "text 1, \"a b c\", is 5 tokens");
  assert_eq!(rows[2].indices, vec![4], "text 2, \"a b\", is 4 tokens");
  assert_eq!(
    rows[0].values,
    vec![1.0],
    "text 0 was the lone second chunk"
  );
  assert_eq!(
    rows[1].values,
    vec![2.0],
    "text 1 shared the first chunk with text 2"
  );
  assert_eq!(
    rows[2].values,
    vec![2.0],
    "text 2 shared the first chunk with text 1"
  );
}

#[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, as `naver/splade-code-06B`'s does.
fn stub_model_no_special_tokens(max_length: usize, batch_size: usize) -> SparseTextEmbedding {
  let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
    .join("tests/fixtures/tiny_tokenizer_no_special_tokens.json");
  SparseTextEmbedding::new(
    Box::new(StubSparseModel),
    TextTokenizer::from_file(&path, max_length, PadToken::Id(0)).unwrap(),
    stub_descriptor(),
    Device::Cpu,
    DType::F32,
    batch_size,
  )
}

/// The stub returns a non-empty row for a zero-length text, so an empty result 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);
  assert_eq!(
    rows[0].indices,
    vec![1],
    "\"a\" is one token with no special tokens added"
  );
  assert_eq!(
    rows[0].values,
    vec![2.0],
    "both rows shared one chunk of two"
  );
  assert_eq!(
    rows[1].indices,
    vec![0],
    "the empty row still reaches StubSparseModel, which echoes its own zero length"
  );
  assert_eq!(
    rows[1].values,
    vec![2.0],
    "both rows shared one chunk of two"
  );
}

/// `SparseTextEmbedding` is not `Debug`, so this matches the `Result` rather than calling `unwrap_err`.
#[test]
fn an_unregistered_id_is_refused_by_name() {
  match SparseTextEmbedding::builder("org/definitely-not-registered").build() {
    Err(CanddingError::UnregisteredModel { id, catalog }) => {
      assert_eq!(id, "org/definitely-not-registered");
      assert_eq!(catalog, Catalog::Sparse);
    }
    Ok(_) => panic!("an unregistered id must be refused"),
    Err(other) => panic!("unexpected error: {other}"),
  }
}

#[test]
fn a_computed_entry_is_refused_before_any_fetch_or_load_is_attempted() {
  match SparseTextEmbedding::builder("Qdrant/bm25").build() {
    Err(CanddingError::ComputedSparseModel { hint }) => {
      assert!(hint.contains("Bm25TextEmbedding"), "{hint}");
    }
    Ok(_) => panic!("a computed entry must be refused, not built"),
    Err(other) => panic!("unexpected error: {other}"),
  }
}

#[test]
fn a_computed_descriptor_is_rejected_and_an_encoded_one_passes_through() {
  let mut computed = stub_descriptor();
  computed.computation = SparseComputation::Computed;
  let err = build::reject_computed_entry(&computed).unwrap_err();
  assert!(
    matches!(err, CanddingError::ComputedSparseModel { .. }),
    "{err}"
  );

  assert!(build::reject_computed_entry(&stub_descriptor()).is_ok());
}