//! Search measurement harness: recall vs. speed for the HNSW index, and
//! end-to-end `hybrid_search` latency as the store grows.
//!
//! Every search-path change should be justified by a before/after run of this
//! binary. It reports, for deterministic synthetic data:
//!
//! * **ANN** — index build time, and for each `ef`: recall@10 against an exact
//!   brute-force scan, queries/second, and p50/p99 latency.
//! * **End to end** — `HDF5Memory`: ingest time, checkpoint time, `open()`
//!   time, the one-off cold index build (first query ever), the first query
//!   after a reopen, and steady-state `hybrid_search` p50/p99 at each size.
//!
//! Data is *clustered* (points = cluster centre + noise, unit-normalised), not
//! uniform: uniform random high-dimensional vectors are nearly equidistant,
//! which makes recall numbers meaningless and is nothing like embeddings.
//!
//! ```text
//! cargo run --release -p clawhdf5-bench --bin search_harness            # 1K, 10K
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --full  # + 100K
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --ann-only --uniform
//! ```

use std::time::{Duration, Instant};

use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
use clawhdf5_ann::{DistanceMetric, HnswIndex};

const DIM: usize = 384;
const K: usize = 10;
const N_QUERIES: usize = 200;
const HNSW_M: usize = 16;
const HNSW_EF_CONSTRUCTION: usize = 64;
const EF_VALUES: [usize; 5] = [16, 32, 64, 128, 256];

// ---------------------------------------------------------------------------
// Deterministic data
// ---------------------------------------------------------------------------

struct Rng(u64);

impl Rng {
    fn next_u64(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }

    /// Uniform in [0, 1).
    fn unit(&mut self) -> f32 {
        (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32
    }

    /// Approximately standard normal (sum of uniforms).
    fn gauss(&mut self) -> f32 {
        let sum: f32 = (0..6).map(|_| self.unit()).sum();
        (sum - 3.0) * std::f32::consts::SQRT_2
    }

    fn below(&mut self, n: usize) -> usize {
        (self.next_u64() % n as u64) as usize
    }
}

fn normalize(v: &mut [f32]) {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        v.iter_mut().for_each(|x| *x /= norm);
    }
}

struct Dataset {
    vectors: Vec<Vec<f32>>,
    queries: Vec<Vec<f32>>,
    /// Cluster id of each vector (used to give records topical text).
    cluster_of: Vec<usize>,
    query_cluster: Vec<usize>,
}

/// `--uniform`: isotropic random unit vectors instead of clusters. Not a
/// realistic workload, but a useful second distribution — a recall problem
/// that appears only on clustered data points at graph connectivity.
static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

fn make_dataset(n: usize, seed: u64) -> Dataset {
    let mut rng = Rng(seed);
    if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) {
        let random_unit = |rng: &mut Rng| {
            let mut v: Vec<f32> = (0..DIM).map(|_| rng.gauss()).collect();
            normalize(&mut v);
            v
        };
        return Dataset {
            vectors: (0..n).map(|_| random_unit(&mut rng)).collect(),
            queries: (0..N_QUERIES).map(|_| random_unit(&mut rng)).collect(),
            cluster_of: vec![0; n],
            query_cluster: vec![0; N_QUERIES],
        };
    }
    let n_clusters = (n / 100).clamp(8, 512);
    let centres: Vec<Vec<f32>> = (0..n_clusters)
        .map(|_| {
            let mut c: Vec<f32> = (0..DIM).map(|_| rng.gauss()).collect();
            normalize(&mut c);
            c
        })
        .collect();
    let point = |rng: &mut Rng, cluster: usize| {
        // Noise comparable to the centre's per-dimension magnitude, so
        // clusters overlap and the nearest neighbours are non-trivial.
        let scale = 0.6 / (DIM as f32).sqrt();
        let mut v: Vec<f32> = centres[cluster]
            .iter()
            .map(|c| c + rng.gauss() * scale)
            .collect();
        normalize(&mut v);
        v
    };
    let mut vectors = Vec::with_capacity(n);
    let mut cluster_of = Vec::with_capacity(n);
    for _ in 0..n {
        let c = rng.below(n_clusters);
        vectors.push(point(&mut rng, c));
        cluster_of.push(c);
    }
    let mut queries = Vec::with_capacity(N_QUERIES);
    let mut query_cluster = Vec::with_capacity(N_QUERIES);
    for _ in 0..N_QUERIES {
        let c = rng.below(n_clusters);
        queries.push(point(&mut rng, c));
        query_cluster.push(c);
    }
    Dataset {
        vectors,
        queries,
        cluster_of,
        query_cluster,
    }
}

const WORDS: &[&str] = &[
    "deploy", "latency", "cache", "schema", "index", "vector", "memory", "agent", "kernel",
    "buffer", "socket", "thread", "tensor", "gradient", "ledger", "invoice", "meeting", "roadmap",
    "customer", "contract", "sensor", "orbit", "protein", "genome", "harbor", "bridge", "engine",
    "battery", "harvest", "weather", "museum", "recipe",
];

/// Text whose vocabulary is biased by cluster, so keyword and vector signals
/// agree the way they do for real embedded text.
fn text_for(cluster: usize, i: usize, rng: &mut Rng) -> String {
    let topic = [
        WORDS[cluster % WORDS.len()],
        WORDS[(cluster / 7 + 3) % WORDS.len()],
    ];
    let mut words = Vec::with_capacity(14);
    for j in 0..14 {
        if j % 3 == 0 {
            words.push(topic[j / 3 % 2]);
        } else {
            words.push(WORDS[rng.below(WORDS.len())]);
        }
    }
    format!("record {i}: {}", words.join(" "))
}

// ---------------------------------------------------------------------------
// Measurement helpers
// ---------------------------------------------------------------------------

fn exact_top_k(vectors: &[Vec<f32>], query: &[f32], k: usize) -> Vec<usize> {
    // Vectors are unit length, so cosine order == dot-product order.
    let mut scored: Vec<(usize, f32)> = vectors
        .iter()
        .enumerate()
        .map(|(i, v)| (i, v.iter().zip(query).map(|(a, b)| a * b).sum()))
        .collect();
    scored.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
    scored.truncate(k);
    scored.into_iter().map(|(i, _)| i).collect()
}

struct Latency {
    p50: Duration,
    p99: Duration,
    qps: f64,
}

fn summarize(mut samples: Vec<Duration>) -> Latency {
    samples.sort();
    let total: Duration = samples.iter().sum();
    let at = |q: f64| samples[((samples.len() - 1) as f64 * q).round() as usize];
    Latency {
        p50: at(0.50),
        p99: at(0.99),
        qps: samples.len() as f64 / total.as_secs_f64(),
    }
}

fn micros(d: Duration) -> f64 {
    d.as_secs_f64() * 1e6
}

fn millis(d: Duration) -> f64 {
    d.as_secs_f64() * 1e3
}

// ---------------------------------------------------------------------------
// ANN: recall vs speed
// ---------------------------------------------------------------------------

fn bench_ann(n: usize, json: &mut Vec<serde_json::Value>) {
    let data = make_dataset(n, 0xA11CE ^ n as u64);
    let truth: Vec<Vec<usize>> = data
        .queries
        .iter()
        .map(|q| exact_top_k(&data.vectors, q, K))
        .collect();

    let started = Instant::now();
    let index = HnswIndex::build_with_metric(
        &data.vectors,
        HNSW_M,
        HNSW_EF_CONSTRUCTION,
        DistanceMetric::Cosine,
    );
    let build = started.elapsed();

    // Exact scan baseline, for scale.
    let exact = summarize(
        data.queries
            .iter()
            .map(|q| {
                let t = Instant::now();
                std::hint::black_box(exact_top_k(&data.vectors, q, K));
                t.elapsed()
            })
            .collect(),
    );

    println!(
        "\n### HNSW, N = {n}, dim = {DIM}, M = {HNSW_M}, ef_construction = {HNSW_EF_CONSTRUCTION}\n"
    );
    println!(
        "build: {:.1} ms ({:.0} vectors/s) · exact scan: {:.0} QPS, p50 {:.0} µs\n",
        millis(build),
        n as f64 / build.as_secs_f64(),
        exact.qps,
        micros(exact.p50)
    );
    println!("| ef | recall@{K} | QPS | p50 µs | p99 µs |");
    println!("|---:|---:|---:|---:|---:|");
    for ef in EF_VALUES {
        let mut hits = 0usize;
        let mut samples = Vec::with_capacity(data.queries.len());
        for (q, want) in data.queries.iter().zip(&truth) {
            let t = Instant::now();
            let got = index.search(q, K, ef);
            samples.push(t.elapsed());
            hits += got.iter().filter(|(id, _)| want.contains(id)).count();
        }
        let recall = hits as f64 / (K * data.queries.len()) as f64;
        let lat = summarize(samples);
        println!(
            "| {ef} | {recall:.4} | {:.0} | {:.0} | {:.0} |",
            lat.qps,
            micros(lat.p50),
            micros(lat.p99)
        );
        json.push(serde_json::json!({
            "bench": "hnsw", "n": n, "ef": ef, "recall_at_10": recall,
            "qps": lat.qps, "p50_us": micros(lat.p50), "p99_us": micros(lat.p99),
            "build_ms": millis(build),
        }));
    }
}

// ---------------------------------------------------------------------------
// End to end: HDF5Memory::hybrid_search
// ---------------------------------------------------------------------------

fn bench_end_to_end(n: usize, json: &mut Vec<serde_json::Value>) {
    let data = make_dataset(n, 0xE2E ^ n as u64);
    let dir = tempfile::TempDir::new().unwrap();
    let path = dir.path().join("store.h5");
    let mut rng = Rng(7);

    let entries: Vec<MemoryEntry> = data
        .vectors
        .iter()
        .enumerate()
        .map(|(i, v)| MemoryEntry {
            chunk: text_for(data.cluster_of[i], i, &mut rng),
            embedding: v.clone(),
            source_channel: "bench".into(),
            timestamp: i as f64,
            session_id: format!("s{}", i % 50),
            tags: format!("t{i}"),
        })
        .collect();
    let query_texts: Vec<String> = data
        .query_cluster
        .iter()
        .enumerate()
        .map(|(i, c)| text_for(*c, i, &mut rng))
        .collect();

    let mut mem = HDF5Memory::create(MemoryConfig::new(path.clone(), "bench", DIM)).unwrap();
    let t = Instant::now();
    mem.save_batch(entries).unwrap();
    let ingest = t.elapsed();
    // The very first query builds the vector and keyword indexes from
    // scratch. It happens once per store, not once per session: the checkpoint
    // below saves the vector index, so a later `open()` reloads it.
    let t = Instant::now();
    std::hint::black_box(mem.hybrid_search(&data.queries[1], &query_texts[1], 0.7, 0.3, K));
    let cold_build = t.elapsed();

    let t = Instant::now();
    mem.flush_wal().unwrap();
    let checkpoint = t.elapsed();
    drop(mem);

    let t = Instant::now();
    let mut mem = HDF5Memory::open(&path).unwrap();
    let open = t.elapsed();

    // The first query after open pays for whatever is rebuilt lazily.
    let t = Instant::now();
    std::hint::black_box(mem.hybrid_search(&data.queries[0], &query_texts[0], 0.7, 0.3, K));
    let first_query = t.elapsed();

    // Fewer steady-state samples at large N: each query is currently O(N).
    let samples_wanted = if n >= 100_000 { 20 } else { N_QUERIES.min(100) };
    let steady = summarize(
        (0..samples_wanted)
            .map(|i| {
                let t = Instant::now();
                std::hint::black_box(mem.hybrid_search(
                    &data.queries[i % N_QUERIES],
                    &query_texts[i % N_QUERIES],
                    0.7,
                    0.3,
                    K,
                ));
                t.elapsed()
            })
            .collect(),
    );

    println!(
        "| {n} | {:.0} | {:.0} | {:.1} | {:.1} | {:.1} | {:.2} | {:.2} | {:.1} |",
        millis(ingest),
        millis(cold_build),
        millis(checkpoint),
        millis(open),
        millis(first_query),
        millis(steady.p50),
        millis(steady.p99),
        steady.qps
    );
    json.push(serde_json::json!({
        "bench": "hybrid_search", "n": n,
        "ingest_ms": millis(ingest), "cold_index_build_ms": millis(cold_build),
        "checkpoint_ms": millis(checkpoint),
        "open_ms": millis(open), "first_query_ms": millis(first_query),
        "p50_ms": millis(steady.p50), "p99_ms": millis(steady.p99), "qps": steady.qps,
    }));
}

// ---------------------------------------------------------------------------
// Fusion study: does capping the keyword candidate pool change the ranking?
// ---------------------------------------------------------------------------

/// `hybrid_search` min-max normalises each signal over the candidates it is
/// given. The vector stage supplies a pool of `max(8k, 64)`; the keyword stage
/// supplies *every* matching record, which is what now dominates query time.
/// This compares the current fusion with one whose keyword stage is capped to
/// a pool, reporting how often the final top-k agree and what each costs.
fn fusion_study(n: usize) {
    use clawhdf5_agent::bm25::BM25Index;
    use clawhdf5_agent::hybrid::merge_vector_keyword;

    let data = make_dataset(n, 0xE2E ^ n as u64);
    let mut rng = Rng(7);
    let texts: Vec<String> = (0..n)
        .map(|i| text_for(data.cluster_of[i], i, &mut rng))
        .collect();
    let query_texts: Vec<String> = data
        .query_cluster
        .iter()
        .enumerate()
        .map(|(i, c)| text_for(*c, i, &mut rng))
        .collect();
    let bm25 = BM25Index::build(&texts, &vec![0u8; n]);
    let index = HnswIndex::build_with_metric(
        &data.vectors,
        HNSW_M,
        HNSW_EF_CONSTRUCTION,
        DistanceMetric::Cosine,
    );

    let vec_pool = (K * 8).max(64);
    println!("\n### Fusion study, N = {n} (k = {K}, weights 0.7 / 0.3, vector pool {vec_pool})\n");
    println!(
        "| keyword pool | top-{K} overlap vs full | identical top-{K} | same #1 | keyword+merge µs |"
    );
    println!("|---:|---:|---:|---:|---:|");

    let fuse = |q: usize, kw_pool: usize| -> (Vec<usize>, Duration) {
        let vec_scores: Vec<(usize, f32)> = index
            .search(&data.queries[q], vec_pool, vec_pool)
            .into_iter()
            .map(|(id, d)| (id, 1.0 - d))
            .collect();
        let t = Instant::now();
        let kw = bm25.search(&query_texts[q], kw_pool);
        let merged = merge_vector_keyword(vec_scores, kw, 0.7, 0.3, K);
        let took = t.elapsed();
        (merged.into_iter().map(|(id, _)| id).collect(), took)
    };

    let full: Vec<(Vec<usize>, Duration)> = (0..N_QUERIES).map(|q| fuse(q, n)).collect();
    let full_time: Duration = full.iter().map(|f| f.1).sum();
    println!(
        "| all ({n}) | 1.0000 | 100.0% | 100.0% | {:.0} |",
        micros(full_time) / N_QUERIES as f64
    );
    for pool in [vec_pool, vec_pool * 4, 1000] {
        if pool >= n {
            continue;
        }
        let (mut overlap, mut identical, mut same_first) = (0usize, 0usize, 0usize);
        let mut time = Duration::ZERO;
        for (q, (want, _)) in full.iter().enumerate() {
            let (got, took) = fuse(q, pool);
            time += took;
            overlap += got.iter().filter(|id| want.contains(id)).count();
            identical += usize::from(&got == want);
            same_first += usize::from(got.first() == want.first());
        }
        println!(
            "| {pool} | {:.4} | {:.1}% | {:.1}% | {:.0} |",
            overlap as f64 / (K * N_QUERIES) as f64,
            100.0 * identical as f64 / N_QUERIES as f64,
            100.0 * same_first as f64 / N_QUERIES as f64,
            micros(time) / N_QUERIES as f64
        );
    }
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let full = args.iter().any(|a| a == "--full");
    let ann_only = args.iter().any(|a| a == "--ann-only");
    if args.iter().any(|a| a == "--fusion-study") {
        for &n in if full {
            &[10_000, 100_000][..]
        } else {
            &[10_000][..]
        } {
            fusion_study(n);
        }
        return;
    }
    if args.iter().any(|a| a == "--uniform") {
        UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed);
        println!("(uniform random data)");
    }
    let json_path = args
        .iter()
        .position(|a| a == "--json")
        .and_then(|i| args.get(i + 1))
        .cloned();
    let sizes: &[usize] = if full {
        &[1_000, 10_000, 100_000]
    } else {
        &[1_000, 10_000]
    };

    if cfg!(debug_assertions) {
        eprintln!("warning: debug build — numbers are meaningless. Use --release.");
    }

    let mut json = Vec::new();
    println!("## Search harness");
    // `--e2e-only` skips the index benchmarks, so the end-to-end section runs
    // in a process that has not already spun up a thread pool.
    if !args.iter().any(|a| a == "--e2e-only") {
        for &n in sizes {
            bench_ann(n, &mut json);
        }
    }

    if ann_only {
        return;
    }
    println!("\n### End to end: `HDF5Memory::hybrid_search` (k = {K}, weights 0.7 / 0.3)\n");
    println!(
        "| N | ingest ms | cold index build ms | checkpoint ms | open ms | first query after open ms | p50 ms | p99 ms | QPS |"
    );
    println!("|---:|---:|---:|---:|---:|---:|---:|---:|---:|");
    for &n in sizes {
        bench_end_to_end(n, &mut json);
    }

    if let Some(path) = json_path {
        std::fs::write(&path, serde_json::to_string_pretty(&json).unwrap()).unwrap();
        eprintln!("wrote {path}");
    }
}