candding 0.1.0

use std::borrow::Cow;

use super::*;
use crate::models::Unembedding;
use crate::tokenizer::PromptTemplate;

const ID: &str = "Qwen/Qwen3-Reranker-0.6B";

#[test]
fn qwen3_reranker_0_6b_is_registered_with_its_own_measured_shape() {
  let descriptor = lookup(ID).expect("Qwen/Qwen3-Reranker-0.6B must be registered");
  assert_eq!(descriptor.id, ID);
  assert_eq!(descriptor.family, Family::Qwen3);
  assert_eq!(
    descriptor.head,
    RerankHead::LogitScore {
      true_token_id: 9693,
      false_token_id: 2152,
      unembedding: Unembedding::Tied,
    },
    "1_LogitScore/config.json, and config.json's tie_word_embeddings: true"
  );
  let json = serde_json::to_value(&descriptor).unwrap();
  assert_eq!(json["head"]["unembedding"], "tied");
  assert_eq!(
    descriptor.score_activation,
    ScoreActivation::Identity,
    "config_sentence_transformers.json's activation_fn"
  );
  assert_eq!(
    descriptor.input,
    RerankInput::Prompt {
      template: PromptTemplate::Qwen3Reranker,
      default_instruction: Cow::Borrowed(
        "Given a web search query, retrieve relevant passages that answer the query"
      ),
    },
    "config_sentence_transformers.json's prompts.query"
  );
  assert_eq!(
    descriptor.default_max_length, 8192,
    "the model card's max_length"
  );
  assert_eq!(descriptor.default_dtype, DType::F32);
  assert_eq!(descriptor.phase, 5);
  assert!(!descriptor.gated);
  assert_eq!(descriptor.license, "Apache-2.0");
}

#[test]
fn qwen3_reranker_0_6b_is_f32_verified_everywhere_and_f16_accelerator_only() {
  assert_accelerator_only_f16_and_bf16(ID, Support::Unsupported);
  let descriptor = lookup(ID).expect("registered");
  let everywhere = BackendSupport {
    cpu: Support::Verified,
    metal: Support::Verified,
    cuda: Support::Verified,
  };
  assert_eq!(descriptor.backends, everywhere);
  assert_eq!(descriptor.dtype_backends.f32, everywhere);
}

fn assert_qwen3_reranker_shape(id: &str, unembedding: Unembedding) -> RerankDescriptor {
  let descriptor = lookup(id).unwrap_or_else(|| panic!("{id} must be registered"));
  assert_eq!(descriptor.id, id);
  assert_eq!(descriptor.family, Family::Qwen3, "{id}");
  assert_eq!(
    descriptor.head,
    RerankHead::LogitScore {
      true_token_id: 9693,
      false_token_id: 2152,
      unembedding,
    },
    "{id}: 1_LogitScore/config.json and config.json's tie_word_embeddings"
  );
  let json = serde_json::to_value(&descriptor).unwrap();
  assert_eq!(json["head"]["unembedding"], unembedding.name(), "{id}");
  assert_eq!(
    descriptor.score_activation,
    ScoreActivation::Identity,
    "{id}"
  );
  assert_eq!(
    descriptor.input,
    RerankInput::Prompt {
      template: PromptTemplate::Qwen3Reranker,
      default_instruction: Cow::Borrowed(
        "Given a web search query, retrieve relevant passages that answer the query"
      ),
    },
    "{id}: config_sentence_transformers.json's prompts.query"
  );
  assert_eq!(
    descriptor.default_max_length, 8192,
    "{id}: the model card's max_length"
  );
  assert_eq!(descriptor.phase, 5, "{id}");
  assert!(!descriptor.gated, "{id}");
  assert_eq!(descriptor.license, "Apache-2.0", "{id}");
  descriptor
}

#[test]
fn qwen3_reranker_4b_is_registered_with_its_own_measured_shape() {
  let descriptor = assert_qwen3_reranker_shape("Qwen/Qwen3-Reranker-4B", Unembedding::Tied);
  assert_eq!(descriptor.default_dtype, DType::BF16);
  let json = serde_json::to_value(&descriptor).unwrap();
  assert_eq!(json["default_dtype"], "bf16");
  assert_eq!(descriptor.size_gb, 8.04, "two shards, 8 043 614 960 bytes");
}

#[test]
fn qwen3_reranker_8b_is_registered_with_its_own_measured_shape() {
  let descriptor = assert_qwen3_reranker_shape("Qwen/Qwen3-Reranker-8B", Unembedding::Untied);
  assert_eq!(descriptor.default_dtype, DType::BF16);
  let json = serde_json::to_value(&descriptor).unwrap();
  assert_eq!(json["default_dtype"], "bf16");
  assert_eq!(
    descriptor.size_gb, 16.38,
    "five shards, 16 377 142 272 bytes"
  );
}

#[test]
fn qwen3_reranker_4b_is_bf16_verified_on_metal_with_f16_unsupported_and_cpu_f32_pending() {
  let id = "Qwen/Qwen3-Reranker-4B";
  let descriptor = lookup(id).expect("registered");
  let row = |cpu, metal, cuda| BackendSupport { cpu, metal, cuda };
  let (pending, verified, untested, unsupported) = (
    Support::Pending,
    Support::Verified,
    Support::Untested,
    Support::Unsupported,
  );
  assert_eq!(descriptor.backends, row(pending, verified, untested));
  assert_eq!(
    descriptor.dtype_backends.f32,
    row(pending, untested, untested)
  );
  assert_eq!(
    descriptor.dtype_backends.f16,
    row(unsupported, unsupported, unsupported)
  );
  assert_eq!(
    descriptor.dtype_backends.bf16,
    row(unsupported, verified, untested)
  );
  assert!(!dtype_needs_accelerator(id, DType::F16));
}

#[test]
fn qwen3_reranker_8b_is_bf16_verified_on_metal_with_metal_f16_and_f32_untested() {
  let id = "Qwen/Qwen3-Reranker-8B";
  let descriptor = lookup(id).expect("registered");
  let row = |cpu, metal, cuda| BackendSupport { cpu, metal, cuda };
  let (verified, untested, unsupported) =
    (Support::Verified, Support::Untested, Support::Unsupported);
  assert_eq!(descriptor.backends, row(untested, verified, untested));
  assert_eq!(
    descriptor.dtype_backends.f32,
    row(untested, untested, untested)
  );
  assert_eq!(
    descriptor.dtype_backends.f16,
    row(unsupported, untested, untested)
  );
  assert_eq!(
    descriptor.dtype_backends.bf16,
    row(unsupported, verified, untested)
  );
  assert!(dtype_needs_accelerator(id, DType::F16));
}