// Meso bench: full Reranker::score_pairs end-to-end against TinyBERT-L-2-v2.
//
// Heritage: @Lampson (1983) "Hints for Computer System Design" - the meso
// bench is the second of three measurement levels (micro / meso / macro).
// Each level catches what the others may miss: micro isolates vtable
// overhead; meso measures the full forward-pass hot path; macro catches
// system-level regressions on real corpora.
//
// This bench loads cross-encoder/ms-marco-TinyBERT-L-2-v2 ONCE in the
// bench-group setup block (not per iteration). Setup time is not counted
// in the iteration measurement; only `score_pairs` wall-clock is sampled.
//
// Model download: weights (~22 MB) are fetched via hf-hub on the first
// run if not already present in `~/.cache/huggingface/hub/`. Pre-warm
// the cache by running:
//   cargo run --release -p ripvec-core --example corpus_bench \
//     --features cpu-accelerate -- tests/corpus/gutenberg \
//     tests/corpus/annotations/gutenberg.json --scope docs --repeats 1
// (see `docs/surgery/backend_trait_microbench.md` Section 6).

#![allow(missing_docs)]

use criterion::{Criterion, criterion_group, criterion_main};
use ripvec_core::rerank::Reranker;

/// HuggingFace repo ID for the default cross-encoder.
///
/// TinyBERT-L-2-v2: 2-layer distilled cross-encoder, ~22 MB, runs at
/// ~33 ms p50 on M2 Max vs 671 ms for the heavier MiniLM-L-12 default
/// (see `rerank.rs` DEFAULT_RERANK_MODEL doc-comment for full sweep table).
const RERANK_MODEL: &str = "cross-encoder/ms-marco-TinyBERT-L-2-v2";

/// Fixed query string used across all iterations for reproducibility.
const QUERY: &str = "what is a static encoder";

/// Number of candidates per batch.
const N_CANDIDATES: usize = 100;

fn bench_full_rerank(c: &mut Criterion) {
    // Load TinyBERT ONCE - not per iteration.
    // Criterion's setup block IS included in wall-clock for the group
    // setup call, but NOT in the per-iteration measurement.
    let reranker = Reranker::from_pretrained(RERANK_MODEL)
        .expect("failed to load TinyBERT-L-2-v2; ensure HuggingFace cache is warm");

    // Build 100 (query, candidate) pairs. Candidates are ~200 tokens
    // each (generated from a fixed template to ensure reproducibility
    // across machines). The template produces a semantically coherent
    // but synthetic sentence padded to approximately 200 tokens by
    // repeating a consistent suffix phrase.
    let candidates: Vec<String> = (0..N_CANDIDATES)
        .map(|i| {
            format!(
                "Static encoders pre-compute token embeddings independently for query and \
                 document without cross-attention, enabling fast approximate nearest-neighbour \
                 retrieval at the cost of cross-token interaction. Candidate {i}. \
                 This document is padded to approximately two hundred tokens to match \
                 real-world rerank input distributions observed on prose and code corpora \
                 in the ripvec benchmark suite. Token budget: static encoder, cross-encoder, \
                 distillation, BERT, TinyBERT, sentence-transformers, ms-marco, passage \
                 retrieval, semantic search, information retrieval, neural reranking."
            )
        })
        .collect();

    // Pre-build pairs as (&str, &str) slices - allocation not part of hot path.
    let pairs: Vec<(&str, &str)> = candidates.iter().map(|c| (QUERY, c.as_str())).collect();

    let mut group = c.benchmark_group("rerank_full");
    // sample_size(50): the TinyBERT forward pass dominates iteration time
    // (~33 ms per batch on M2 Max); 50 samples deliver stable p50/p95
    // estimates. The default 100 samples would double bench runtime with
    // no meaningful accuracy gain. Criterion will emit a warning about
    // reduced sample size; this is expected and intentional.
    // Per @Box (1979): aim for usefully informative measurement, not
    // maximally precise measurement.
    group.sample_size(50);
    group.bench_function("score_batch_100", |b| {
        b.iter(|| {
            let scores = reranker
                .score_pairs(std::hint::black_box(&pairs))
                .expect("rerank failed");
            std::hint::black_box(scores);
        });
    });
    group.finish();
}

criterion_group!(benches, bench_full_rerank);
criterion_main!(benches);
