use std::{
    hint::black_box,
    io::{Cursor, Read},
};

#[cfg(feature = "bench-allocations")]
use std::fmt::Write as _;
#[cfg(not(feature = "bench-allocations"))]
use std::time::Duration;

#[cfg(not(feature = "bench-allocations"))]
use criterion::{BenchmarkId, SamplingMode, Throughput};
use criterion::{Criterion, criterion_group, criterion_main};
use kvdb_overlay::{Database, DatabaseOverlay, Value};

#[cfg(feature = "fjall-backend")]
const BACKEND: &str = "fjall";
#[cfg(feature = "sled-backend")]
const BACKEND: &str = "sled";

const TREE: &str = "value_reads";
const KEY: &[u8] = b"value";
const SIZES: [(&str, usize); 7] = [
    ("64B", 64),
    ("4KiB", 4096),
    ("1MiB", 1024 * 1024),
    ("32MiB", 32 * 1024 * 1024),
    ("128MiB", 128 * 1024 * 1024),
    ("256MiB", 256 * 1024 * 1024),
    ("512MiB", 512 * 1024 * 1024),
];

// Both decoders use the same generic Read implementation, with no decode heap allocation.
fn decode(mut reader: impl Read, words: usize) -> u64 {
    let mut checksum = 0u64;
    let mut word = [0u8; 8];
    for _ in 0..words {
        reader.read_exact(&mut word).unwrap();
        checksum = checksum.wrapping_add(u64::from_le_bytes(word));
    }
    checksum
}

type ReadValue = fn(&DatabaseOverlay, usize) -> u64;
const READERS: [(&str, ReadValue); 4] = [
    ("vec_cursor", |overlay, words| {
        let value = overlay.get(TREE, KEY).unwrap().unwrap().into_vec();
        decode(Cursor::new(value), words)
    }),
    ("vec_slice", |overlay, words| {
        let value = overlay.get(TREE, KEY).unwrap().unwrap().into_vec();
        decode(value.as_slice(), words)
    }),
    ("handle_slice", |overlay, words| {
        let value = overlay.get(TREE, KEY).unwrap().unwrap();
        decode(value.as_ref(), words)
    }),
    ("handle_cursor", |overlay, words| {
        let value = overlay.get(TREE, KEY).unwrap().unwrap();
        decode(Cursor::new(value), words)
    }),
];

fn value_reads(_criterion: &mut Criterion) {
    let selection = std::env::var("VALUE_READ_SIZE").unwrap_or_else(|_| "all".into());
    assert!(
        selection == "all"
            || selection == "large"
            || SIZES.iter().any(|(name, _)| *name == selection),
        "VALUE_READ_SIZE must be all, large, 64B, 4KiB, 1MiB, 32MiB, 128MiB, 256MiB, or 512MiB"
    );
    #[cfg(feature = "bench-allocations")]
    let mut report = format!(
        "# Value Read Allocations ({BACKEND})\n\n\
         Generated by `VALUE_READ_SIZE={selection} cargo bench --bench value_reads --no-default-features \
         --features {BACKEND}-backend,bench-allocations`.\n\n\
         Warm repeated get + decode + drop; foreground thread only. \
         Totals are cumulative allocations, not live/peak memory or RSS. \
         Backend worker allocations are excluded. No Criterion timings in this run.\n\n\
         | backend | source | size_bytes | decode | approach | operations | allocations_total | bytes_total | allocations_per_op | bytes_per_op |\n\
         | --- | --- | ---: | --- | --- | ---: | ---: | ---: | ---: | ---: |\n"
    );

    for source in ["backend_fallback", "overlay_cache"] {
        for (size_name, size) in SIZES {
            let large = size >= 32 * 1024 * 1024;
            if selection != "all" && selection != size_name && !(selection == "large" && large) {
                continue;
            }
            // Bound fixed validation/allocation work for large payloads.
            let repeats = (100 * 1024 * 1024 / size).clamp(3, 100);
            // Build deterministic, non-uniform words and independent expected checksums.
            let mut payload = Vec::with_capacity(size);
            let mut expected_full = 0u64;
            let mut expected_header = 0u64;
            for index in 0..size / 8 {
                let word = (index as u64 + 1).wrapping_mul(0x9e3779b97f4a7c15);
                if index == 0 {
                    expected_header = word;
                }
                expected_full = expected_full.wrapping_add(word);
                payload.extend_from_slice(&word.to_le_bytes());
            }

            // Keep the directory alive until all database/overlay handles are dropped.
            let directory = tempfile::tempdir().unwrap();
            let database = Database::open_default(directory.path()).unwrap();
            let tree = database
                .open_tree(
                    TREE,
                    #[cfg(feature = "fjall-backend")]
                    || {
                        let options = kvdb_overlay::fjall::KeyspaceCreateOptions::default();
                        if large {
                            // Keep the warm-handle experiment on the same storage path.
                            options.max_memtable_size(1024 * 1024 * 1024)
                        } else {
                            options
                        }
                    },
                )
                .unwrap();
            tree.insert(KEY, &payload).unwrap();
            database.flush_default_mode().unwrap();
            let mut overlay = DatabaseOverlay::new(&database, vec![]).unwrap();
            overlay.open_tree_default(TREE, false).unwrap();
            if source == "overlay_cache" {
                overlay.insert(TREE, KEY, &payload).unwrap();
            }
            assert_eq!(
                matches!(overlay.get(TREE, KEY).unwrap().unwrap(), Value::Cached(_)),
                source == "overlay_cache"
            );
            drop(payload);

            #[cfg(feature = "fjall-backend")]
            let assert_resident = || {
                assert_eq!(tree.tree().sealed_memtable_count(), 0);
                assert_eq!(tree.tree().table_count(), 0);
            };

            #[cfg(not(feature = "bench-allocations"))]
            let mut group =
                _criterion.benchmark_group(format!("value_reads/{BACKEND}/{source}/{size_name}"));
            #[cfg(not(feature = "bench-allocations"))]
            {
                group.sample_size(25);
                group.measurement_time(Duration::from_secs(1));
                group.warm_up_time(Duration::from_millis(300));
                group.throughput(Throughput::Elements(1));
                if large {
                    group.sampling_mode(SamplingMode::Flat);
                    group.measurement_time(Duration::from_secs(3));
                }
            }

            for (decode_name, words, expected) in [
                ("header_8B", 1, expected_header),
                ("full_payload", size / 8, expected_full),
            ] {
                for (approach, read_value) in READERS {
                    // Validate every approach and warm this exact path before measuring.
                    for _ in 0..repeats {
                        assert_eq!(read_value(black_box(&overlay), black_box(words)), expected);
                    }
                    #[cfg(feature = "fjall-backend")]
                    assert_resident();

                    #[cfg(feature = "bench-allocations")]
                    {
                        let allocations = allocation_counter::measure(|| {
                            for _ in 0..repeats {
                                black_box(read_value(black_box(&overlay), black_box(words)));
                            }
                        });
                        writeln!(
                            report,
                            "| {BACKEND} | {source} | {size} | {decode_name} | {approach} | {repeats} | {} | {} | {:.2} | {:.2} |",
                            allocations.count_total,
                            allocations.bytes_total,
                            allocations.count_total as f64 / repeats as f64,
                            allocations.bytes_total as f64 / repeats as f64,
                        ).unwrap();
                    }

                    #[cfg(not(feature = "bench-allocations"))]
                    group.bench_function(BenchmarkId::new(decode_name, approach), |bencher| {
                        bencher.iter(|| {
                            // The owned Vec/handle is dropped inside read_value, before return.
                            black_box(read_value(black_box(&overlay), black_box(words)))
                        });
                    });
                    #[cfg(feature = "fjall-backend")]
                    assert_resident();
                }
            }
            #[cfg(not(feature = "bench-allocations"))]
            group.finish();
        }
    }

    #[cfg(feature = "bench-allocations")]
    {
        let directory = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("target/criterion");
        std::fs::create_dir_all(&directory).unwrap();
        let suffix = if selection == "all" {
            String::new()
        } else {
            format!("-{selection}")
        };
        let path = directory.join(format!("value-allocations-{BACKEND}{suffix}.md"));
        std::fs::write(&path, &report).unwrap();
        print!("{report}");
        eprintln!("Allocation report: {}", path.display());
    }
}

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