use std::hint::black_box;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};

use codspeed_criterion_compat::{BenchmarkId, Criterion, criterion_group, criterion_main};
use lru::LruCache;

mod common;
use common::{build_compressible_payload, deflate_compress};

const BLOCK_SIZE: usize = 4096;
const SHARED_REPEATS: usize = 64;
const MIXED_DISTINCT: usize = 10;
const MIXED_REPEATS: usize = 10;

fn decompress(compressed: &[u8]) -> Vec<u8> {
    use std::io::Read;
    let mut out = Vec::with_capacity(BLOCK_SIZE);
    let mut d = flate2::read::DeflateDecoder::new(compressed);
    d.read_to_end(&mut out).expect("inflate");
    out
}

struct BenchLru {
    inner: Mutex<LruCache<u64, Arc<Vec<u8>>>>,
}

impl BenchLru {
    fn with_capacity(n: usize) -> Self {
        let cap = NonZeroUsize::new(n.max(1)).expect("cap >= 1");
        Self {
            inner: Mutex::new(LruCache::new(cap)),
        }
    }

    fn get_or_inflate(&self, key: u64, compressed: &[u8]) -> Arc<Vec<u8>> {
        if let Some(hit) = self.inner.lock().expect("bench lock").get(&key).cloned() {
            return hit;
        }
        let decompressed = Arc::new(decompress(compressed));
        let _ = self
            .inner
            .lock()
            .expect("bench lock")
            .put(key, Arc::clone(&decompressed));
        decompressed
    }
}

fn build_block(seed: u8) -> Vec<u8> {
    let mut raw = build_compressible_payload(BLOCK_SIZE);
    raw[0] ^= seed;
    deflate_compress(&raw)
}

fn build_unique_blocks(n: usize) -> Vec<Vec<u8>> {
    (0..n)
        .map(|i| build_block(u8::try_from(i % 256).unwrap()))
        .collect()
}

fn unique_blocks(c: &mut Criterion) {
    const N: usize = 64;
    let blocks = build_unique_blocks(N);
    let mut g = c.benchmark_group("unique_blocks");
    g.sample_size(20);

    g.bench_with_input(BenchmarkId::new("baseline", N), &(), |b, ()| {
        b.iter(|| {
            let mut total = 0usize;
            for blob in &blocks {
                total += black_box(decompress(blob)).len();
            }
            black_box(total);
        });
    });

    g.bench_with_input(BenchmarkId::new("lru", N), &(), |b, ()| {
        b.iter(|| {
            let cache = BenchLru::with_capacity(N);
            let mut total = 0usize;
            for (i, blob) in blocks.iter().enumerate() {
                total += cache.get_or_inflate(i as u64, blob).len();
            }
            black_box(total);
        });
    });

    g.finish();
}

fn shared_block(c: &mut Criterion) {
    let blob = build_block(0);
    let mut g = c.benchmark_group("shared_block");
    g.sample_size(20);

    g.bench_with_input(
        BenchmarkId::new("baseline", SHARED_REPEATS),
        &(),
        |b, ()| {
            b.iter(|| {
                let mut total = 0usize;
                for _ in 0..SHARED_REPEATS {
                    total += black_box(decompress(&blob)).len();
                }
                black_box(total);
            });
        },
    );

    g.bench_with_input(BenchmarkId::new("lru", SHARED_REPEATS), &(), |b, ()| {
        b.iter(|| {
            let cache = BenchLru::with_capacity(8);
            let mut total = 0usize;
            for _ in 0..SHARED_REPEATS {
                total += cache.get_or_inflate(0, &blob).len();
            }
            black_box(total);
        });
    });

    g.finish();
}

fn mixed_block(c: &mut Criterion) {
    let blocks = build_unique_blocks(MIXED_DISTINCT);
    let access_order: Vec<usize> = {
        let mut out = Vec::with_capacity(MIXED_DISTINCT * MIXED_REPEATS);
        let stride = 7;
        let mut idx = 0usize;
        for _ in 0..(MIXED_DISTINCT * MIXED_REPEATS) {
            out.push(idx % MIXED_DISTINCT);
            idx += stride;
        }
        out
    };

    let mut g = c.benchmark_group("mixed_block");
    g.sample_size(20);

    g.bench_with_input(
        BenchmarkId::new("baseline", access_order.len()),
        &(),
        |b, ()| {
            b.iter(|| {
                let mut total = 0usize;
                for &i in &access_order {
                    total += black_box(decompress(&blocks[i])).len();
                }
                black_box(total);
            });
        },
    );

    g.bench_with_input(BenchmarkId::new("lru", access_order.len()), &(), |b, ()| {
        b.iter(|| {
            let cache = BenchLru::with_capacity(MIXED_DISTINCT);
            let mut total = 0usize;
            for &i in &access_order {
                total += cache.get_or_inflate(i as u64, &blocks[i]).len();
            }
            black_box(total);
        });
    });

    g.finish();
}

criterion_group!(benches, unique_blocks, shared_block, mixed_block);
criterion_main!(benches);
