candding 0.1.0

use std::collections::HashMap;
use std::fs;

use candle_core::DType;
use candle_core::Device;
use tempfile::TempDir;

use super::*;
use crate::MultiVectorTextEmbedding;
use crate::SparseTextEmbedding;
use crate::TextCrossEncoder;
use crate::TextEmbedding;
use crate::embed::synthetic;
use crate::pooling::Pooling;
use crate::registry;
use crate::registry::Family;
use crate::registry::ModelDescriptor;

#[test]
fn weights_that_need_exactly_what_is_available_fit() {
  assert!(fits(0, 0));
  assert!(fits(1 << 30, 1 << 30));
  assert!(fits(u64::MAX, u64::MAX));
}

#[test]
fn one_byte_over_what_is_available_does_not_fit() {
  assert!(!fits(1, 0));
  assert!(!fits((1 << 30) + 1, 1 << 30));
  assert!(!fits(u64::MAX, u64::MAX - 1));
}

#[test]
fn a_refusal_carries_the_model_the_backend_and_both_sizes() {
  assert!(verdict("m", Backend::Metal, 10, Some(10)).is_ok());
  assert!(
    verdict("m", Backend::Metal, 11, None).is_ok(),
    "no figure is no limit"
  );
  match verdict("m", Backend::Metal, 11, Some(10)) {
    Err(CanddingError::InsufficientMemory {
      model,
      backend,
      needed_bytes,
      available_bytes,
    }) => assert_eq!(
      (model.as_str(), backend, needed_bytes, available_bytes),
      ("m", Backend::Metal, 11, 10)
    ),
    other => panic!("{other:?}"),
  }
}

/// 2^50 F32 elements, 4 PiB, declared with none of their bytes: a check that came after reading them would fail there.
const DECLARED_BYTES: u64 = 4 << 50;

fn header_only_checkpoint() -> TempDir {
  let config = synthetic::bert_config(16);
  let dir = synthetic::checkpoint(&config, "tiny_tokenizer.json", &HashMap::new());
  let header =
    br#"{"w":{"dtype":"F32","shape":[1125899906842624],"data_offsets":[0,4503599627370496]}}"#;
  let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
  bytes.extend_from_slice(header);
  fs::write(dir.path().join("model.safetensors"), bytes).unwrap();
  dir
}

fn assert_refused_on_the_cpu(refusal: Option<CanddingError>, needed: u64) {
  match refusal {
    Some(CanddingError::InsufficientMemory {
      backend: Backend::Cpu,
      needed_bytes,
      ..
    }) => assert_eq!(needed_bytes, needed),
    other => panic!("{other:?}"),
  }
}

#[cfg(any(target_os = "linux", target_os = "macos"))]
#[test]
fn every_builder_refuses_weights_the_cpu_cannot_hold_before_reading_them() {
  let dir = header_only_checkpoint();
  let spec = dir.path().to_str().unwrap();
  let cpu = || Device::Cpu;
  let refusals = [
    TextEmbedding::builder(spec).device(cpu()).build().err(),
    SparseTextEmbedding::builder(spec)
      .descriptor(registry::sparse::lookup("prithivida/Splade_PP_en_v1").unwrap())
      .device(cpu())
      .build()
      .err(),
    MultiVectorTextEmbedding::builder(spec)
      .descriptor(registry::multi_vector::lookup("answerdotai/answerai-colbert-small-v1").unwrap())
      .device(cpu())
      .build()
      .err(),
    TextCrossEncoder::builder(spec)
      .descriptor(registry::rerank::lookup("cross-encoder/ms-marco-MiniLM-L6-v2").unwrap())
      .device(cpu())
      .build()
      .err(),
  ];
  for refusal in refusals {
    assert_refused_on_the_cpu(refusal, DECLARED_BYTES);
  }
}

/// No snapshot of the id exists, so a check placed after `hub::resolve` would reach the network instead.
#[cfg(any(target_os = "linux", target_os = "macos"))]
#[test]
fn a_model_not_on_disk_is_refused_from_its_registry_size_before_the_download() {
  let id = "rust-dd/candding-fit-check-never-published";
  for (stored, needed) in [
    (DType::F32, 1_000_000_000_000_000_000),
    (DType::BF16, 2_000_000_000_000_000_000),
  ] {
    let mut descriptor = ModelDescriptor::new(id, Family::Bert, 4, Pooling::Mean);
    descriptor.size_gb = 1e9;
    descriptor.default_dtype = stored;
    let refusal = TextEmbedding::builder(id)
      .descriptor(descriptor)
      .device(Device::Cpu)
      .dtype(DType::F32)
      .build()
      .err();
    assert_refused_on_the_cpu(refusal, needed);
  }
}

#[test]
fn a_snapshot_on_disk_is_sized_from_its_headers_not_its_registry_size() {
  let config = synthetic::bert_config(16);
  let dir = synthetic::checkpoint(
    &config,
    "tiny_tokenizer.json",
    &synthetic::zero_bert("", 16),
  );
  let spec = dir.path().to_str().unwrap();
  let mut descriptor = ModelDescriptor::new(spec, Family::Bert, synthetic::HIDDEN, Pooling::Mean);
  descriptor.size_gb = 1e9;
  let built = TextEmbedding::builder(spec)
    .descriptor(descriptor)
    .device(Device::Cpu)
    .build();
  assert!(built.is_ok(), "{:?}", built.err());
}