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

use dsa::{FastStore, MainStore, RandomAccessMemory};

// ── reproduce the trait and PhysAddr here, or import from your crate ─────────
type PhysAddr = u32;

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

/// Cycles through a small set of addresses — stays hot in L1/L2.
/// Simulates a tight inner loop hitting the same working set repeatedly.
fn sequential_addrs(n: usize, max_addr: u32) -> Vec<PhysAddr> {
    (0..n).map(|i| ((i as u32 * 4) & (max_addr - 1))).collect()
}

/// LCG pseudo-random addresses across the full address space.
/// Simulates worst-case cache behaviour.
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
}

/// Walks addresses page by page — tests page-boundary crossing behaviour.
fn page_stride_addrs(n: usize) -> Vec<PhysAddr> {
    (0..n).map(|i| (i as u32 * 4096) & 0x00FF_FFFF).collect()
}

// ── generic benchmark functions ───────────────────────────────────────────────

fn bench_read_byte<M: RandomAccessMemory>(
    group: &mut BenchmarkGroup<WallTime>,
    name: &str,
    mem: &M,
    addrs: &[PhysAddr],
) {
    group.throughput(Throughput::Elements(addrs.len() as u64));
    group.bench_function(name, |b| {
        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 bench_write_byte<M: RandomAccessMemory>(
    group: &mut BenchmarkGroup<WallTime>,
    name: &str,
    mem: &mut M,
    addrs: &[PhysAddr],
) {
    group.throughput(Throughput::Elements(addrs.len() as u64));
    group.bench_function(name, |b| {
        b.iter(|| {
            for (i, &addr) in addrs.iter().enumerate() {
                mem.write_byte(black_box(addr), black_box(i as u8));
            }
        })
    });
}

fn bench_read_word<M: RandomAccessMemory>(
    group: &mut BenchmarkGroup<WallTime>,
    name: &str,
    mem: &M,
    addrs: &[PhysAddr],
) {
    group.throughput(Throughput::Elements(addrs.len() as u64));
    group.bench_function(name, |b| {
        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 bench_write_word<M: RandomAccessMemory>(
    group: &mut BenchmarkGroup<WallTime>,
    name: &str,
    mem: &mut M,
    addrs: &[PhysAddr],
) {
    group.throughput(Throughput::Elements(addrs.len() as u64));
    group.bench_function(name, |b| {
        b.iter(|| {
            for (i, &addr) in addrs.iter().enumerate() {
                mem.write_word(black_box(addr), black_box(i as u32));
            }
        })
    });
}

fn bench_read_page<M: RandomAccessMemory>(
    group: &mut BenchmarkGroup<WallTime>,
    name: &str,
    mem: &M,
    addrs: &[PhysAddr],
) {
    group.throughput(Throughput::Bytes(addrs.len() as u64 * 4096));
    group.bench_function(name, |b| {
        b.iter(|| {
            let mut sum: u8 = 0;
            for &addr in addrs {
                let page = mem.read_page(black_box(addr));
                sum = sum.wrapping_add(page[0]);
            }
            black_box(sum)
        })
    });
}

fn bench_write_page<M: RandomAccessMemory>(
    group: &mut BenchmarkGroup<WallTime>,
    name: &str,
    mem: &mut M,
    addrs: &[PhysAddr],
) {
    let page_data = [0xABu8; 4096];
    group.throughput(Throughput::Bytes(addrs.len() as u64 * 4096));
    group.bench_function(name, |b| {
        b.iter(|| {
            for &addr in addrs {
                mem.write_page(black_box(addr), black_box(&page_data));
            }
        })
    });
}

// ── run all access patterns for a given implementation ────────────────────────

fn bench_implementation<M: RandomAccessMemory>(
    c: &mut Criterion,
    impl_name: &str,
    mut mem: M,
    max_addr: u32,
) {
    const N: usize = 64;

    let seq_addrs = sequential_addrs(N, max_addr);
    let rand_addrs = random_addrs(N, max_addr);
    let page_stride = page_stride_addrs(N);

    // ── read_byte ────────────────────────────────────────────────────────────
    {
        let mut g = c.benchmark_group(format!("{}/read_byte", impl_name));
        g.measurement_time(Duration::from_secs(3));
        bench_read_byte(&mut g, "sequential", &mem, &seq_addrs);
        bench_read_byte(&mut g, "random", &mem, &rand_addrs);
        g.finish();
    }

    // ── write_byte ───────────────────────────────────────────────────────────
    {
        let mut g = c.benchmark_group(format!("{}/write_byte", impl_name));
        g.measurement_time(Duration::from_secs(3));
        bench_write_byte(&mut g, "sequential", &mut mem, &seq_addrs);
        bench_write_byte(&mut g, "random", &mut mem, &rand_addrs);
        g.finish();
    }

    // ── read_word ────────────────────────────────────────────────────────────
    {
        let mut g = c.benchmark_group(format!("{}/read_word", impl_name));
        g.measurement_time(Duration::from_secs(3));
        bench_read_word(&mut g, "sequential", &mem, &seq_addrs);
        bench_read_word(&mut g, "random", &mem, &rand_addrs);
        g.finish();
    }

    // ── write_word ───────────────────────────────────────────────────────────
    {
        let mut g = c.benchmark_group(format!("{}/write_word", impl_name));
        g.measurement_time(Duration::from_secs(3));
        bench_write_word(&mut g, "sequential", &mut mem, &seq_addrs);
        bench_write_word(&mut g, "random", &mut mem, &rand_addrs);
        g.finish();
    }

    // ── read_page ────────────────────────────────────────────────────────────
    {
        let mut g = c.benchmark_group(format!("{}/read_page", impl_name));
        g.measurement_time(Duration::from_secs(3));
        bench_read_page(&mut g, "sequential", &mem, &seq_addrs);
        bench_read_page(&mut g, "random", &mem, &rand_addrs);
        bench_read_page(&mut g, "page_stride", &mem, &page_stride);
        g.finish();
    }

    // ── write_page ───────────────────────────────────────────────────────────
    {
        let mut g = c.benchmark_group(format!("{}/write_page", impl_name));
        g.measurement_time(Duration::from_secs(3));
        bench_write_page(&mut g, "sequential", &mut mem, &seq_addrs);
        bench_write_page(&mut g, "random", &mut mem, &rand_addrs);
        bench_write_page(&mut g, "page_stride", &mut mem, &page_stride);
        g.finish();
    }
}

// ── wire up your implementations here ────────────────────────────────────────
//
// Replace these stubs with your actual types. Each call to bench_implementation
// runs the full suite and labels it separately in the HTML report.
//
// Example:
//   fn benchmarks(c: &mut Criterion) {
//       bench_implementation(c, "MainStore",  MainStore::new(),  0x00FF_FFFF);
//       bench_implementation(c, "FxHashMap",  HashMapMem::new(), 0x00FF_FFFF);
//   }
//
fn benchmarks(c: &mut Criterion) {
    // TODO: replace with your implementations
    bench_implementation(c, "MainStore", MainStore::new(), 0x00FF_FFFF);
    bench_implementation(c, "FastStore", FastStore::new(), 0x00FF_FFFF);

    let _ = c;
}

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