use criterion::{
    BenchmarkGroup, BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main,
    measurement::WallTime,
};
use std::time::Duration;

use dsa::{Page, RandomAccessMemory};

type PhysAddr = u32;

// ── address generators ────────────────────────────────────────────────────────

fn sequential_addrs(n: usize, max_addr: u32) -> Vec<PhysAddr> {
    (0..n).map(|i| ((i as u32 * 4) & (max_addr - 1))).collect()
}

fn random_addrs(n: usize, max_addr: u32) -> Vec<PhysAddr> {
    let mut addrs = Vec::with_capacity(n);
    let mut x: u32 = 0xdeadbeef;
    for _ in 0..n {
        x = x.wrapping_mul(1664525).wrapping_add(1013904223);
        addrs.push((x & (max_addr - 1)) & !3);
    }
    addrs
}

fn page_stride_addrs(n: usize) -> Vec<PhysAddr> {
    (0..n).map(|i| (i as u32 * 4096) & 0x00FF_FFFF).collect()
}

fn random_page_addrs(n: usize, max_addr: u32) -> Vec<PhysAddr> {
    let mut addrs = Vec::with_capacity(n);
    let mut x: u32 = 0xc0ffee42;
    for _ in 0..n {
        x = x.wrapping_mul(1664525).wrapping_add(1013904223);
        addrs.push((x & (max_addr - 1)) & !0xFFF);
    }
    addrs
}

// ── runners ───────────────────────────────────────────────────────────────────

fn run_read_byte(
    group: &mut BenchmarkGroup<WallTime>,
    addrs: &[PhysAddr],
    implementations: &mut [(&str, Box<dyn RandomAccessMemory>)],
) {
    group.throughput(Throughput::Elements(addrs.len() as u64));
    for (name, mem) in implementations.iter() {
        group.bench_with_input(BenchmarkId::new(*name, ""), addrs, |b, addrs| {
            b.iter(|| {
                let mut sum: u32 = 0;
                for &addr in addrs {
                    sum = sum.wrapping_add(mem.read_byte(black_box(addr)) as u32);
                }
                black_box(sum)
            })
        });
    }
}

fn run_write_byte(
    group: &mut BenchmarkGroup<WallTime>,
    addrs: &[PhysAddr],
    implementations: &mut [(&str, Box<dyn RandomAccessMemory>)],
) {
    group.throughput(Throughput::Elements(addrs.len() as u64));
    for (name, mem) in implementations.iter_mut() {
        group.bench_with_input(BenchmarkId::new(*name, ""), addrs, |b, addrs| {
            b.iter(|| {
                for (i, &addr) in addrs.iter().enumerate() {
                    mem.write_byte(black_box(addr), black_box(i as u8));
                }
            })
        });
    }
}

fn run_read_word(
    group: &mut BenchmarkGroup<WallTime>,
    addrs: &[PhysAddr],
    implementations: &mut [(&str, Box<dyn RandomAccessMemory>)],
) {
    group.throughput(Throughput::Elements(addrs.len() as u64));
    for (name, mem) in implementations.iter() {
        group.bench_with_input(BenchmarkId::new(*name, ""), addrs, |b, addrs| {
            b.iter(|| {
                let mut sum: u32 = 0;
                for &addr in addrs {
                    sum = sum.wrapping_add(mem.read_word(black_box(addr)));
                }
                black_box(sum)
            })
        });
    }
}

fn run_write_word(
    group: &mut BenchmarkGroup<WallTime>,
    addrs: &[PhysAddr],
    implementations: &mut [(&str, Box<dyn RandomAccessMemory>)],
) {
    group.throughput(Throughput::Elements(addrs.len() as u64));
    for (name, mem) in implementations.iter_mut() {
        group.bench_with_input(BenchmarkId::new(*name, ""), addrs, |b, addrs| {
            b.iter(|| {
                for (i, &addr) in addrs.iter().enumerate() {
                    mem.write_word(black_box(addr), black_box(i as u32));
                }
            })
        });
    }
}

fn run_read_page(
    group: &mut BenchmarkGroup<WallTime>,
    addrs: &[PhysAddr],
    implementations: &mut [(&str, Box<dyn RandomAccessMemory>)],
) {
    group.throughput(Throughput::Bytes(addrs.len() as u64 * 4096));
    for (name, mem) in implementations.iter() {
        group.bench_with_input(BenchmarkId::new(*name, ""), addrs, |b, addrs| {
            b.iter(|| {
                let mut sum: u64 = 0;
                for &addr in addrs {
                    let page = mem.read_page(black_box(addr));
                    for chunk in page.chunks_exact(8) {
                        sum = sum.wrapping_add(u64::from_le_bytes(chunk.try_into().unwrap()));
                    }
                }
                black_box(sum)
            })
        });
    }
}

fn run_write_page(
    group: &mut BenchmarkGroup<WallTime>,
    addrs: &[PhysAddr],
    implementations: &mut [(&str, Box<dyn RandomAccessMemory>)],
) {
    let page_data = Page::from([0xABu8; 4096]);
    group.throughput(Throughput::Bytes(addrs.len() as u64 * 4096));
    for (name, mem) in implementations.iter_mut() {
        group.bench_with_input(BenchmarkId::new(*name, ""), addrs, |b, addrs| {
            b.iter(|| {
                for &addr in addrs {
                    mem.write_page(black_box(addr), black_box(&page_data));
                }
            })
        });
    }
}

// ── entry point ───────────────────────────────────────────────────────────────

fn benchmarks(c: &mut Criterion) {
    const N: usize = 4096;
    const MAX_ADDR: u32 = 0x00FF_FFFF;

    let seq_addrs = sequential_addrs(N, MAX_ADDR);
    let rand_addrs = random_addrs(N, MAX_ADDR);
    let page_addrs_seq = page_stride_addrs(N);
    let page_addrs_rand = random_page_addrs(N, MAX_ADDR);

    let mut implementations: Vec<(&str, Box<dyn RandomAccessMemory>)> = vec![
        #[cfg(feature = "mainstore-hashmap")]
        ("HashStore", Box::new(dsa::HashStore::new())),
        #[cfg(feature = "mainstore-arraymap")]
        ("ArrayStore", Box::new(dsa::ArrayStore::new())),
        #[cfg(feature = "mainstore-stackarray")]
        ("StackArrayStore", Box::new(dsa::StackArrayStore::new())),
        #[cfg(feature = "mainstore-prealloc")]
        ("PreAllocStore", Box::new(dsa::PreAllocStore::new())),
        #[cfg(feature = "mainstore-bulkalloc")]
        ("BulkAllocStore", Box::new(dsa::BulkAllocStore::new())),
    ];

    macro_rules! group {
        ($name:expr, $secs:expr, $runner:ident, $addrs:expr) => {{
            let mut g = c.benchmark_group($name);
            g.measurement_time(Duration::from_secs($secs));
            $runner(&mut g, $addrs, &mut implementations);
            g.finish();
        }};
    }

    group!("write_word/random", 30, run_write_word, &rand_addrs);
    group!("write_byte/random", 30, run_write_byte, &rand_addrs);

    group!("read_byte/sequential", 30, run_read_byte, &seq_addrs);
    group!("read_byte/random", 30, run_read_byte, &rand_addrs);

    group!("read_word/sequential", 30, run_read_word, &seq_addrs);
    group!("read_word/random", 30, run_read_word, &rand_addrs);

    group!("write_byte/sequential", 30, run_write_byte, &seq_addrs);
    group!("write_word/sequential", 30, run_write_word, &seq_addrs);

    group!("read_page/sequential", 30, run_read_page, &page_addrs_seq);
    group!("read_page/random", 30, run_read_page, &page_addrs_rand);

    group!("write_page/sequential", 30, run_write_page, &page_addrs_seq);
    group!("write_page/random", 30, run_write_page, &page_addrs_rand);
}

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