//! Multi-modal memory search benchmarks (`clawhdf5_agent::multimodal`).
//!
//! Covers `MultiModalStore::search_cross_modal` (every embedding of every
//! record, whatever its modality) and, for comparison,
//! `MultiModalStore::search_by_modality` restricted to one modality.
//!
//! Corpus: N records (1K and 10K), each carrying two 384-dim embeddings —
//! a text embedding of its caption plus one embedding of its primary modality,
//! cycling Image / Audio / Video — so a cross-modal query scores 2N vectors.
//! All data comes from a fixed-seed LCG, so every run sees the same corpus.
//!
//! Run: `cargo bench -p clawhdf5-agent --bench multimodal_bench`

use std::collections::HashMap;

use clawhdf5_agent::multimodal::{
    MediaRef, ModalEmbedding, Modality, MultiModalRecord, MultiModalStore,
};
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};

// ---------------------------------------------------------------------------
// Simple deterministic PRNG (LCG), same as the other agent benches
// ---------------------------------------------------------------------------

struct Rng(u32);

impl Rng {
    fn new(seed: u32) -> Self {
        Self(seed)
    }
    fn next_u32(&mut self) -> u32 {
        self.0 = self.0.wrapping_mul(1103515245).wrapping_add(12345);
        self.0 >> 16
    }
    fn next_f32(&mut self) -> f32 {
        self.next_u32() as f32 / 65536.0 - 0.5
    }
}

fn make_vec(rng: &mut Rng, dim: usize) -> Vec<f32> {
    (0..dim).map(|_| rng.next_f32()).collect()
}

// ---------------------------------------------------------------------------
// Corpus
// ---------------------------------------------------------------------------

const DIM: usize = 384;
const K: usize = 10;

const MEDIA: [(Modality, &str, &str); 3] = [
    (Modality::Image, "image/png", "clip-vit-base"),
    (Modality::Audio, "audio/wav", "clap-base"),
    (Modality::Video, "video/mp4", "xclip-base"),
];

fn build_store(n: usize, seed: u32) -> MultiModalStore {
    let mut rng = Rng::new(seed);
    let mut store = MultiModalStore::new();
    for i in 0..n {
        let (modality, mime, model) = &MEDIA[i % MEDIA.len()];
        let embeddings = vec![
            ModalEmbedding::new(Modality::Text, make_vec(&mut rng, DIM), "minilm-l6"),
            ModalEmbedding::new(modality.clone(), make_vec(&mut rng, DIM), *model),
        ];
        store.add_record(MultiModalRecord {
            id: 0,
            primary_modality: modality.clone(),
            text_content: Some(format!("{modality} memory {i}")),
            media_ref: Some(MediaRef::path(format!("/media/{i}"), *mime)),
            embeddings,
            observation: None,
            timestamp: 1_700_000_000.0 + i as f64,
            metadata: HashMap::new(),
        });
    }
    store
}

// ---------------------------------------------------------------------------
// Benchmarks
// ---------------------------------------------------------------------------

fn multimodal_search_benches(c: &mut Criterion) {
    let query = make_vec(&mut Rng::new(99), DIM);

    let mut group = c.benchmark_group("multimodal_search");
    group.sample_size(50);

    for (label, n) in [("1k", 1_000usize), ("10k", 10_000)] {
        let store = build_store(n, 42);
        assert_eq!(store.count(), n);

        group.bench_with_input(BenchmarkId::new("cross_modal", label), &n, |b, _| {
            b.iter(|| store.search_cross_modal(&query, K));
        });

        group.bench_with_input(BenchmarkId::new("by_modality_image", label), &n, |b, _| {
            b.iter(|| store.search_by_modality(&Modality::Image, &query, K));
        });
    }

    group.finish();
}

criterion_group!(multimodal_benches, multimodal_search_benches);
criterion_main!(multimodal_benches);