//! Criterion bench for TierStore promotion + demotion (wave 4).

use std::hint::black_box;

use criterion::{criterion_group, criterion_main, Criterion};
use jouleclaw_loader_gguf::synthetic::{synthesize_llama_gguf, SyntheticConfig};
use jouleclaw_loader_gguf::tier_store::{StoreTensorMeta, TierConfig, TierStore};
use jouleclaw_loader_gguf::{read_gguf_file, GgufBuffer};
use std::io::Write;
use std::path::PathBuf;

fn write_synth_gguf() -> PathBuf {
    let bytes = synthesize_llama_gguf(&SyntheticConfig::default());
    let dir = std::env::temp_dir();
    let pid = std::process::id();
    let path = dir.join(format!("jouleclaw-bench-tier-{}.gguf", pid));
    let mut f = std::fs::File::create(&path).unwrap();
    f.write_all(&bytes).unwrap();
    f.sync_all().unwrap();
    path
}

fn setup_store() -> (TierStore, PathBuf, Vec<StoreTensorMeta>) {
    let path = write_synth_gguf();
    let model = read_gguf_file(&path).unwrap();
    let data_offset = match &model.buf {
        GgufBuffer::Mapped { data_offset, .. } => *data_offset as u64,
        _ => panic!(),
    };
    let mmap_len = match &model.buf {
        GgufBuffer::Mapped { mmap, .. } => mmap.len() as u64,
        _ => panic!(),
    };
    let data_end = mmap_len - data_offset;
    let mut sorted: Vec<usize> = (0..model.tensors.len()).collect();
    sorted.sort_by_key(|&i| model.tensors[i].offset);
    let mut metas = Vec::with_capacity(model.tensors.len());
    for (rank, &idx) in sorted.iter().enumerate() {
        let t = &model.tensors[idx];
        let next_off = sorted
            .get(rank + 1)
            .map(|&j| model.tensors[j].offset)
            .unwrap_or(data_end);
        metas.push(StoreTensorMeta {
            name: t.name.clone(),
            offset_in_data: t.offset,
            size_bytes: next_off.saturating_sub(t.offset),
        });
    }
    let store = TierStore::new(TierConfig::default());
    store.register_model("m1", &path, data_offset).unwrap();
    store.register_tensors("m1", &metas);
    (store, path, metas)
}

fn bench_promote_warm(c: &mut Criterion) {
    let (store, path, metas) = setup_store();
    c.bench_function("tier_store_promote_to_warm_first_tensor", |b| {
        let first = &metas[0];
        b.iter(|| {
            // Re-promote ad nauseam — first call is Cold→Warm,
            // subsequent are idempotent Warm→Warm. Bench
            // measures the steady-state hot-path latency.
            let _ = store.promote_to_warm(
                black_box("m1"),
                black_box(first),
                black_box(1),
            );
        });
    });
    let _ = std::fs::remove_file(path);
}

fn bench_pin_hot(c: &mut Criterion) {
    let (store, path, metas) = setup_store();
    c.bench_function("tier_store_pin_hot_first_tensor", |b| {
        let first = &metas[0];
        let mut tick: u64 = 1;
        b.iter(|| {
            let _ = store.pin_hot(
                black_box("m1"),
                black_box(first),
                black_box(tick),
            );
            tick = tick.wrapping_add(1);
        });
    });
    let _ = std::fs::remove_file(path);
}

fn bench_bytes_borrow(c: &mut Criterion) {
    let (store, path, metas) = setup_store();
    let first = &metas[0];
    store.promote_to_warm("m1", first, 1).unwrap();
    c.bench_function("tier_store_bytes_borrow", |b| {
        b.iter(|| {
            let len = store
                .bytes(black_box("m1"), black_box(first), |bytes| bytes.len())
                .unwrap();
            black_box(len);
        });
    });
    let _ = std::fs::remove_file(path);
}

criterion_group!(benches, bench_promote_warm, bench_pin_hot, bench_bytes_borrow);
criterion_main!(benches);