use std::collections::HashMap;

use bytes::Bytes;
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use wgcloak::{
    protocol::{
        direct::{parse_packet_bounded, DataPacket},
        replay::ReplayWindow,
    },
    transport::{http3, tls_tcp::FrameCodec},
};

const SECRET: &[u8] = b"benchmark-secret-with-at-least-32-bytes";
const SESSION_KEY: [u8; 32] = [0x42; 32];
const SESSION_ID: u64 = 0x0102_0304_0506_0708;

fn bench_direct_protocol(c: &mut Criterion) {
    let mut group = c.benchmark_group("direct_protocol");
    for size in [64usize, 512, 1200] {
        let payload = Bytes::from(vec![0x5a; size]);
        let packet = DataPacket::authenticated(&SESSION_KEY, SESSION_ID, 7, payload);
        group.throughput(Throughput::Bytes(size as u64));
        group.bench_with_input(
            BenchmarkId::new("encode_parse_mac", size),
            &packet,
            |b, packet| {
                b.iter(|| {
                    let wire = packet.encode().expect("benchmark packet encodes");
                    let parsed = parse_packet_bounded(black_box(&wire), size)
                        .expect("benchmark packet parses");
                    let parsed = match parsed {
                        wgcloak::protocol::Packet::Data(data) => data,
                        _ => panic!("benchmark packet changed type"),
                    };
                    black_box(parsed.verify_authenticator(&SESSION_KEY));
                });
            },
        );
    }
    group.finish();
}

fn bench_tls_frames(c: &mut Criterion) {
    let codec = FrameCodec::new(1200).expect("valid frame maximum");
    let first = vec![0x11; 512];
    let second = vec![0x22; 512];
    let single = codec.encode(&first).expect("frame encodes");
    let mut coalesced = single.clone();
    coalesced.extend_from_slice(&codec.encode(&second).expect("frame encodes"));
    let mut group = c.benchmark_group("tls_frames");
    group.throughput(Throughput::Bytes(single.len() as u64));
    group.bench_function("encode", |b| {
        b.iter(|| black_box(codec.encode(black_box(&first)).expect("frame encodes")))
    });
    group.bench_function("decode_split", |b| {
        b.iter(|| {
            let mut decoder = FrameCodec::new(1200).expect("valid frame maximum");
            let midpoint = single.len() / 2;
            black_box(
                decoder
                    .feed(&single[..midpoint])
                    .expect("first half accepted"),
            );
            black_box(
                decoder
                    .feed(&single[midpoint..])
                    .expect("second half accepted"),
            );
        });
    });
    group.bench_function("decode_coalesced", |b| {
        b.iter(|| {
            let mut decoder = FrameCodec::new(1200).expect("valid frame maximum");
            black_box(decoder.feed(&coalesced).expect("coalesced frames accepted"));
        });
    });
    group.finish();
}

fn bench_http3(c: &mut Criterion) {
    let challenge = http3::AuthChallenge { bytes: [0x33; 32] };
    let nonce = [0x44; 16];
    let client_id = "benchmark-client";
    let response = http3::encode_auth_response(&challenge, client_id, &nonce, SECRET)
        .expect("auth response encodes");
    let mut clients = HashMap::new();
    clients.insert(client_id.to_owned(), SECRET.to_vec());
    let datagram_payload = vec![0x55; 1200];
    let mut group = c.benchmark_group("http3_primitives");
    group.throughput(Throughput::Bytes(datagram_payload.len() as u64));
    group.bench_function("datagram_encode_decode", |b| {
        b.iter(|| {
            let encoded = http3::encode_http_datagram(
                http3::CONNECT_UDP_CONTEXT_ID,
                black_box(&datagram_payload),
                1200,
            )
            .expect("HTTP Datagram encodes");
            black_box(
                http3::decode_http_datagram(&encoded, http3::CONNECT_UDP_CONTEXT_ID, 1200)
                    .expect("HTTP Datagram decodes"),
            );
        });
    });
    group.bench_function("capsule_encode_decode", |b| {
        b.iter(|| {
            let encoded =
                http3::encode_capsule(0x2a, black_box(&challenge.bytes)).expect("capsule encodes");
            black_box(http3::decode_capsule(&encoded).expect("capsule decodes"));
        });
    });
    group.bench_function("auth_response_verify", |b| {
        b.iter(|| {
            black_box(
                http3::verify_auth_response(&challenge, &response, &clients)
                    .expect("auth response verifies"),
            );
        });
    });
    group.finish();
}

fn bench_replay(c: &mut Criterion) {
    let sequence: Vec<u32> = (1..=1024).collect();
    c.bench_function("replay_window_sequential_1024", |b| {
        b.iter(|| {
            let mut window = ReplayWindow::new();
            for sequence in &sequence {
                black_box(window.check_and_record(*sequence));
            }
        });
    });
    c.bench_function("replay_window_duplicate_check", |b| {
        b.iter(|| {
            let mut window = ReplayWindow::new();
            for sequence in 1..=1024 {
                window.check_and_record(sequence);
            }
            black_box(window.check_and_record(512));
        });
    });
}

criterion_group!(
    protocol_transports,
    bench_direct_protocol,
    bench_tls_frames,
    bench_http3,
    bench_replay
);
criterion_main!(protocol_transports);