use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use btest_rs::client::scan_status_message;
use btest_rs::protocol::STATUS_MSG_TYPE;

/// Naive O(n) byte-by-byte scan — the old implementation.
fn naive_scan(buf: &[u8]) -> Option<u8> {
    const STATUS_MSG_SIZE: usize = 12;
    if buf.len() < STATUS_MSG_SIZE {
        return None;
    }
    for i in 0..=(buf.len() - STATUS_MSG_SIZE) {
        if buf[i] == STATUS_MSG_TYPE && buf[i + 1] >= 0x80 {
            return Some((buf[i + 1] & 0x7F).min(100));
        }
    }
    None
}

fn make_buffer(size: usize, status_at: Option<usize>) -> Vec<u8> {
    let mut buf = vec![0u8; size];
    if let Some(pos) = status_at {
        buf[pos] = STATUS_MSG_TYPE;
        buf[pos + 1] = 0x80 | 50; // CPU = 50%
    }
    buf
}

fn bench_scan_all_zeros(c: &mut Criterion) {
    let mut group = c.benchmark_group("tcp_rx_scan_all_zeros");
    for size in [4096, 65536, 262144] {
        let buf = make_buffer(size, None);
        group.throughput(Throughput::Bytes(size as u64));
        group.bench_with_input(BenchmarkId::new("naive", size), &buf, |b, buf| {
            b.iter(|| black_box(naive_scan(black_box(buf))))
        });
        group.bench_with_input(BenchmarkId::new("memchr", size), &buf, |b, buf| {
            b.iter(|| black_box(scan_status_message(black_box(&[]), black_box(buf))))
        });
    }
    group.finish();
}

fn bench_scan_status_at_start(c: &mut Criterion) {
    let mut group = c.benchmark_group("tcp_rx_scan_status_at_start");
    for size in [4096, 65536, 262144] {
        let buf = make_buffer(size, Some(0));
        group.throughput(Throughput::Bytes(size as u64));
        group.bench_with_input(BenchmarkId::new("naive", size), &buf, |b, buf| {
            b.iter(|| black_box(naive_scan(black_box(buf))))
        });
        group.bench_with_input(BenchmarkId::new("memchr", size), &buf, |b, buf| {
            b.iter(|| black_box(scan_status_message(black_box(&[]), black_box(buf))))
        });
    }
    group.finish();
}

fn bench_scan_status_at_end(c: &mut Criterion) {
    let mut group = c.benchmark_group("tcp_rx_scan_status_at_end");
    for size in [4096, 65536, 262144] {
        let buf = make_buffer(size, Some(size - 12));
        group.throughput(Throughput::Bytes(size as u64));
        group.bench_with_input(BenchmarkId::new("naive", size), &buf, |b, buf| {
            b.iter(|| black_box(naive_scan(black_box(buf))))
        });
        group.bench_with_input(BenchmarkId::new("memchr", size), &buf, |b, buf| {
            b.iter(|| black_box(scan_status_message(black_box(&[]), black_box(buf))))
        });
    }
    group.finish();
}

fn bench_scan_split_message(c: &mut Criterion) {
    // Simulate a status message split across two reads:
    // carry has first 5 bytes, buf has remaining 7 bytes
    let mut carry = vec![0u8; 5];
    carry[0] = STATUS_MSG_TYPE;
    carry[1] = 0x80 | 75;
    let buf = vec![0u8; 7];

    c.bench_function("scan_split_5_7", |b| {
        b.iter(|| black_box(scan_status_message(black_box(&carry), black_box(&buf))))
    });

    // Split with 2 bytes in carry (status type + cpu byte), 10 in buf
    let carry_2 = vec![STATUS_MSG_TYPE, 0x80 | 33];
    let buf_10 = vec![0u8; 10];
    c.bench_function("scan_split_2_10", |b| {
        b.iter(|| black_box(scan_status_message(black_box(&carry_2), black_box(&buf_10))))
    });
}

criterion_group!(
    tcp_rx_scan_benches,
    bench_scan_all_zeros,
    bench_scan_status_at_start,
    bench_scan_status_at_end,
    bench_scan_split_message
);
criterion_main!(tcp_rx_scan_benches);