//! Benchmarks for prop_store vs JSON sidecar property storage.
//!
//! Run with: cargo bench --bench prop_store_benchmarks

use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use geographdb_core::storage::prop_store::{
    build_prop_store, lookup, read_prop_store, write_prop_store, NodeProperties, PropertyValue,
};
use std::collections::BTreeMap;
use tempfile::tempdir;

fn sample_props(seed: u64) -> NodeProperties {
    let mut props = BTreeMap::new();
    props.insert("kind".to_string(), PropertyValue::String("hub".to_string()));
    props.insert(
        "active".to_string(),
        PropertyValue::Bool(seed.is_multiple_of(2)),
    );
    props.insert("priority".to_string(), PropertyValue::U64(seed));
    props.insert("score".to_string(), PropertyValue::F64((seed as f64) * 0.1));
    props.insert(
        "embedding".to_string(),
        PropertyValue::F64Array(vec![
            seed as f64 * 0.01,
            seed as f64 * 0.02,
            seed as f64 * 0.03,
            seed as f64 * 0.04,
            seed as f64 * 0.05,
            seed as f64 * 0.06,
            seed as f64 * 0.07,
            seed as f64 * 0.08,
        ]),
    );
    props.insert(
        "label".to_string(),
        PropertyValue::String(format!("node_{}", seed)),
    );
    props
}

fn sample_json_props(seed: u64) -> BTreeMap<String, serde_json::Value> {
    let mut props = BTreeMap::new();
    props.insert(
        "kind".to_string(),
        serde_json::Value::String("hub".to_string()),
    );
    props.insert(
        "active".to_string(),
        serde_json::Value::Bool(seed.is_multiple_of(2)),
    );
    props.insert(
        "priority".to_string(),
        serde_json::Value::Number(serde_json::Number::from(seed)),
    );
    props.insert(
        "score".to_string(),
        serde_json::Value::Number(
            serde_json::Number::from_f64((seed as f64) * 0.1)
                .unwrap_or_else(|| serde_json::Number::from(0)),
        ),
    );
    props.insert(
        "embedding".to_string(),
        serde_json::Value::Array(
            [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08]
                .iter()
                .map(|&base| {
                    serde_json::Value::Number(
                        serde_json::Number::from_f64(seed as f64 * base)
                            .unwrap_or_else(|| serde_json::Number::from(0)),
                    )
                })
                .collect(),
        ),
    );
    props.insert(
        "label".to_string(),
        serde_json::Value::String(format!("node_{}", seed)),
    );
    props
}

fn bench_prop_store_build(c: &mut Criterion) {
    let mut group = c.benchmark_group("prop_store_build");
    for count in [1_000, 10_000, 100_000].iter() {
        group.throughput(Throughput::Elements(*count as u64));
        group.bench_with_input(BenchmarkId::from_parameter(count), count, |b, &count| {
            let nodes: Vec<(u64, NodeProperties)> = (0..count)
                .map(|i| (i as u64, sample_props(i as u64)))
                .collect();
            b.iter(|| {
                let store = build_prop_store(black_box(nodes.clone()));
                black_box(store);
            });
        });
    }
    group.finish();
}

fn bench_json_encode(c: &mut Criterion) {
    let mut group = c.benchmark_group("json_encode");
    for count in [1_000, 10_000, 100_000].iter() {
        group.throughput(Throughput::Elements(*count as u64));
        group.bench_with_input(BenchmarkId::from_parameter(count), count, |b, &count| {
            let props: Vec<BTreeMap<String, serde_json::Value>> =
                (0..count).map(|i| sample_json_props(i as u64)).collect();
            let ids: Vec<u64> = (0..count).map(|i| i as u64).collect();
            b.iter(|| {
                let sidecar = serde_json::json!({
                    "id_map": black_box(ids.clone()),
                    "properties": black_box(props.clone()),
                });
                let _bytes = serde_json::to_vec(&sidecar).unwrap();
                black_box(_bytes);
            });
        });
    }
    group.finish();
}

fn bench_prop_store_write(c: &mut Criterion) {
    let mut group = c.benchmark_group("prop_store_write");
    for count in [1_000, 10_000, 100_000].iter() {
        group.throughput(Throughput::Elements(*count as u64));
        group.bench_with_input(BenchmarkId::from_parameter(count), count, |b, &count| {
            let nodes: Vec<(u64, NodeProperties)> = (0..count)
                .map(|i| (i as u64, sample_props(i as u64)))
                .collect();
            let store = build_prop_store(nodes);
            let tmp = tempdir().unwrap();
            let path = tmp.path().join("test.prop");
            b.iter(|| {
                write_prop_store(black_box(&store), black_box(&path)).unwrap();
            });
        });
    }
    group.finish();
}

fn bench_json_write(c: &mut Criterion) {
    let mut group = c.benchmark_group("json_write");
    for count in [1_000, 10_000, 100_000].iter() {
        group.throughput(Throughput::Elements(*count as u64));
        group.bench_with_input(BenchmarkId::from_parameter(count), count, |b, &count| {
            let props: Vec<BTreeMap<String, serde_json::Value>> =
                (0..count).map(|i| sample_json_props(i as u64)).collect();
            let ids: Vec<u64> = (0..count).map(|i| i as u64).collect();
            let sidecar = serde_json::json!({
                "id_map": ids,
                "properties": props,
            });
            let bytes = serde_json::to_vec(&sidecar).unwrap();
            let tmp = tempdir().unwrap();
            let path = tmp.path().join("test.json");
            b.iter(|| {
                std::fs::write(black_box(&path), black_box(&bytes)).unwrap();
            });
        });
    }
    group.finish();
}

fn bench_prop_store_read(c: &mut Criterion) {
    let mut group = c.benchmark_group("prop_store_read");
    for count in [1_000, 10_000, 100_000].iter() {
        group.throughput(Throughput::Elements(*count as u64));
        group.bench_with_input(BenchmarkId::from_parameter(count), count, |b, &count| {
            let nodes: Vec<(u64, NodeProperties)> = (0..count)
                .map(|i| (i as u64, sample_props(i as u64)))
                .collect();
            let store = build_prop_store(nodes);
            let tmp = tempdir().unwrap();
            let path = tmp.path().join("test.prop");
            write_prop_store(&store, &path).unwrap();
            b.iter(|| {
                let loaded = read_prop_store(black_box(&path)).unwrap();
                black_box(loaded);
            });
        });
    }
    group.finish();
}

fn bench_json_read(c: &mut Criterion) {
    let mut group = c.benchmark_group("json_read");
    for count in [1_000, 10_000, 100_000].iter() {
        group.throughput(Throughput::Elements(*count as u64));
        group.bench_with_input(BenchmarkId::from_parameter(count), count, |b, &count| {
            let props: Vec<BTreeMap<String, serde_json::Value>> =
                (0..count).map(|i| sample_json_props(i as u64)).collect();
            let ids: Vec<u64> = (0..count).map(|i| i as u64).collect();
            let sidecar = serde_json::json!({
                "id_map": ids,
                "properties": props,
            });
            let bytes = serde_json::to_vec(&sidecar).unwrap();
            let tmp = tempdir().unwrap();
            let path = tmp.path().join("test.json");
            std::fs::write(&path, &bytes).unwrap();
            b.iter(|| {
                let data = std::fs::read(black_box(&path)).unwrap();
                let _parsed: serde_json::Value = serde_json::from_slice(&data).unwrap();
                black_box(_parsed);
            });
        });
    }
    group.finish();
}

fn bench_prop_store_lookup(c: &mut Criterion) {
    let mut group = c.benchmark_group("prop_store_lookup");
    for count in [1_000, 10_000, 100_000].iter() {
        group.throughput(Throughput::Elements(*count as u64));
        group.bench_with_input(BenchmarkId::from_parameter(count), count, |b, &count| {
            let nodes: Vec<(u64, NodeProperties)> = (0..count)
                .map(|i| (i as u64, sample_props(i as u64)))
                .collect();
            let store = build_prop_store(nodes);
            let targets: Vec<u64> = (0..count).map(|i| i as u64).collect();
            b.iter(|| {
                for target in &targets {
                    let result = lookup(black_box(&store), black_box(*target)).unwrap();
                    black_box(result);
                }
            });
        });
    }
    group.finish();
}

fn bench_json_lookup(c: &mut Criterion) {
    let mut group = c.benchmark_group("json_lookup");
    for count in [1_000, 10_000, 100_000].iter() {
        group.throughput(Throughput::Elements(*count as u64));
        group.bench_with_input(BenchmarkId::from_parameter(count), count, |b, &count| {
            let props: Vec<BTreeMap<String, serde_json::Value>> =
                (0..count).map(|i| sample_json_props(i as u64)).collect();
            let ids: Vec<u64> = (0..count).map(|i| i as u64).collect();
            let sidecar = serde_json::json!({
                "id_map": ids,
                "properties": props,
            });
            let data = serde_json::to_vec(&sidecar).unwrap();
            let parsed: serde_json::Value = serde_json::from_slice(&data).unwrap();
            let targets: Vec<u64> = (0..count).map(|i| i as u64).collect();
            b.iter(|| {
                let parsed_ref = black_box(&parsed);
                let arr = parsed_ref["properties"].as_array().unwrap();
                for target in &targets {
                    let _props = &arr[*target as usize];
                    black_box(_props);
                }
            });
        });
    }
    group.finish();
}

fn bench_file_size_comparison(c: &mut Criterion) {
    let mut group = c.benchmark_group("file_size_bytes");
    for count in [1_000, 10_000, 100_000].iter() {
        let nodes: Vec<(u64, NodeProperties)> = (0..*count)
            .map(|i| (i as u64, sample_props(i as u64)))
            .collect();
        let store = build_prop_store(nodes);
        let tmp = tempdir().unwrap();
        let prop_path = tmp.path().join("bench.prop");
        write_prop_store(&store, &prop_path).unwrap();
        let prop_size = std::fs::metadata(&prop_path).unwrap().len();

        let json_props: Vec<BTreeMap<String, serde_json::Value>> =
            (0..*count).map(|i| sample_json_props(i as u64)).collect();
        let ids: Vec<u64> = (0..*count).map(|i| i as u64).collect();
        let sidecar = serde_json::json!({"id_map": ids, "properties": json_props});
        let json_bytes = serde_json::to_vec(&sidecar).unwrap();
        let json_path = tmp.path().join("bench.json");
        std::fs::write(&json_path, &json_bytes).unwrap();
        let json_size = std::fs::metadata(&json_path).unwrap().len();

        eprintln!(
            "file_size/{}: prop_store={} bytes, json={} bytes, ratio={:.2}",
            count,
            prop_size,
            json_size,
            json_size as f64 / prop_size as f64
        );

        group.bench_function(BenchmarkId::new("prop_store", count), |b| {
            b.iter(|| black_box(prop_size))
        });
        group.bench_function(BenchmarkId::new("json", count), |b| {
            b.iter(|| black_box(json_size))
        });
    }
    group.finish();
}

criterion_group!(
    benches,
    bench_prop_store_build,
    bench_json_encode,
    bench_prop_store_write,
    bench_json_write,
    bench_prop_store_read,
    bench_json_read,
    bench_prop_store_lookup,
    bench_json_lookup,
    bench_file_size_comparison,
);
criterion_main!(benches);
