use bytes::{Bytes, BytesMut};
use criterion::{Criterion, criterion_group, criterion_main};
use tokio_util::codec::{Decoder, Encoder};

use fafb_q::codec::{Frame, MessageCodec};
use fafb_q::types::{ApiKey, GroupName, Header, MessageId, NamespaceId, TopicName, TopicOffset};

fn make_codec() -> MessageCodec {
    MessageCodec::new(1_048_576)
}

fn make_large_codec() -> MessageCodec {
    MessageCodec::new(2_097_152) // 2MB — for benchmarks with 1MB payloads
}

fn codec_encode(c: &mut Criterion) {
    let mut group = c.benchmark_group("codec_encode");

    group.bench_function("ping", |b| {
        let mut codec = make_codec();
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(64);
            codec.encode(Frame::Ping, &mut buf).unwrap();
        });
    });

    group.bench_function("connect", |b| {
        let mut codec = make_codec();
        let frame = Frame::Connect {
            version: 1,
            namespace: NamespaceId("production-namespace".into()),
            api_key: ApiKey("sk-abcdefghijklmnopqrstuvwxyz1234567890".into()),
            capabilities: 0,
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(128);
            codec.encode(frame.clone(), &mut buf).unwrap();
        });
    });

    group.bench_function("subscribe", |b| {
        let mut codec = make_codec();
        let frame = Frame::Subscribe {
            topic: TopicName("events.user.created".into()),
            group: Some(GroupName("worker-pool-alpha".into())),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(64);
            codec.encode(frame.clone(), &mut buf).unwrap();
        });
    });

    group.bench_function("publish_64b", |b| {
        let mut codec = make_codec();
        let frame = Frame::Publish {
            topic: TopicName("events".into()),
            headers: vec![],
            ttl_ms: 0,
            payload: Bytes::from(vec![0xABu8; 64]),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(128);
            codec.encode(frame.clone(), &mut buf).unwrap();
        });
    });

    group.bench_function("publish_1kb", |b| {
        let mut codec = make_codec();
        let frame = Frame::Publish {
            topic: TopicName("events".into()),
            headers: vec![],
            ttl_ms: 0,
            payload: Bytes::from(vec![0xABu8; 1024]),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(1200);
            codec.encode(frame.clone(), &mut buf).unwrap();
        });
    });

    group.bench_function("publish_64kb", |b| {
        let mut codec = make_codec();
        let frame = Frame::Publish {
            topic: TopicName("events".into()),
            headers: vec![],
            ttl_ms: 0,
            payload: Bytes::from(vec![0xABu8; 65536]),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(66000);
            codec.encode(frame.clone(), &mut buf).unwrap();
        });
    });

    group.bench_function("publish_1mb", |b| {
        let mut codec = make_codec();
        let frame = Frame::Publish {
            topic: TopicName("events".into()),
            headers: vec![],
            ttl_ms: 0,
            payload: Bytes::from(vec![0xABu8; 1_048_576]),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(1_050_000);
            codec.encode(frame.clone(), &mut buf).unwrap();
        });
    });

    group.bench_function("deliver_1kb", |b| {
        let mut codec = make_codec();
        let frame = Frame::Deliver {
            message_id: MessageId(123456),
            offset: TopicOffset(999),
            timestamp_millis: 1_700_000_000_000,
            ttl_ms: 0,
            topic: TopicName("events.user.created".into()),
            headers: vec![],
            payload: Bytes::from(vec![0xCDu8; 1024]),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(1200);
            codec.encode(frame.clone(), &mut buf).unwrap();
        });
    });

    group.finish();
}

fn codec_decode(c: &mut Criterion) {
    let mut group = c.benchmark_group("codec_decode");

    let encode = |frame: Frame| -> Vec<u8> {
        let mut codec = make_codec();
        let mut buf = BytesMut::new();
        codec.encode(frame, &mut buf).unwrap();
        buf.to_vec()
    };

    let pong_bytes = encode(Frame::Pong);
    group.bench_function("pong", |b| {
        let mut codec = make_codec();
        b.iter(|| {
            let mut buf = BytesMut::from(&pong_bytes[..]);
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    let pub_64_bytes = encode(Frame::Publish {
        topic: TopicName("events".into()),
        headers: vec![],
        ttl_ms: 0,
        payload: Bytes::from(vec![0xABu8; 64]),
    });
    group.bench_function("publish_64b", |b| {
        let mut codec = make_codec();
        b.iter(|| {
            let mut buf = BytesMut::from(&pub_64_bytes[..]);
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    let pub_1k_bytes = encode(Frame::Publish {
        topic: TopicName("events".into()),
        headers: vec![],
        ttl_ms: 0,
        payload: Bytes::from(vec![0xABu8; 1024]),
    });
    group.bench_function("publish_1kb", |b| {
        let mut codec = make_codec();
        b.iter(|| {
            let mut buf = BytesMut::from(&pub_1k_bytes[..]);
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    let pub_64k_bytes = encode(Frame::Publish {
        topic: TopicName("events".into()),
        headers: vec![],
        ttl_ms: 0,
        payload: Bytes::from(vec![0xABu8; 65536]),
    });
    group.bench_function("publish_64kb", |b| {
        let mut codec = make_codec();
        b.iter(|| {
            let mut buf = BytesMut::from(&pub_64k_bytes[..]);
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    let pub_1m_bytes = encode(Frame::Publish {
        topic: TopicName("events".into()),
        headers: vec![],
        ttl_ms: 0,
        payload: Bytes::from(vec![0xABu8; 1_048_576]),
    });
    group.bench_function("publish_1mb", |b| {
        let mut codec = make_large_codec();
        b.iter(|| {
            let mut buf = BytesMut::from(&pub_1m_bytes[..]);
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    let deliver_1k_bytes = encode(Frame::Deliver {
        message_id: MessageId(42),
        offset: TopicOffset(100),
        timestamp_millis: 1_700_000_000_000,
        ttl_ms: 0,
        topic: TopicName("events.user.created".into()),
        headers: vec![],
        payload: Bytes::from(vec![0xCDu8; 1024]),
    });
    group.bench_function("deliver_1kb", |b| {
        let mut codec = make_codec();
        b.iter(|| {
            let mut buf = BytesMut::from(&deliver_1k_bytes[..]);
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    group.bench_function("partial_buffer_none", |b| {
        let mut codec = make_codec();
        // Only 3 bytes — always returns Ok(None)
        let partial = vec![0u8; 3];
        b.iter(|| {
            let mut buf = BytesMut::from(&partial[..]);
            assert!(codec.decode(&mut buf).unwrap().is_none());
        });
    });

    group.finish();
}

fn codec_roundtrip(c: &mut Criterion) {
    let mut group = c.benchmark_group("codec_roundtrip");

    group.bench_function("connect", |b| {
        let mut codec = make_codec();
        let frame = Frame::Connect {
            version: 1,
            namespace: NamespaceId("production".into()),
            api_key: ApiKey("sk-abcdefghijklmnopqrstuvwxyz".into()),
            capabilities: 0,
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(128);
            codec.encode(frame.clone(), &mut buf).unwrap();
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    group.bench_function("subscribe", |b| {
        let mut codec = make_codec();
        let frame = Frame::Subscribe {
            topic: TopicName("events.user.created".into()),
            group: Some(GroupName("workers".into())),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(64);
            codec.encode(frame.clone(), &mut buf).unwrap();
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    group.bench_function("publish_1kb", |b| {
        let mut codec = make_codec();
        let frame = Frame::Publish {
            topic: TopicName("events".into()),
            headers: vec![],
            ttl_ms: 0,
            payload: Bytes::from(vec![0xABu8; 1024]),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(1200);
            codec.encode(frame.clone(), &mut buf).unwrap();
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    group.bench_function("publish_1mb", |b| {
        let mut codec = make_large_codec();
        let frame = Frame::Publish {
            topic: TopicName("events".into()),
            headers: vec![],
            ttl_ms: 0,
            payload: Bytes::from(vec![0xABu8; 1_048_576]),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(1_050_000);
            codec.encode(frame.clone(), &mut buf).unwrap();
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    group.bench_function("deliver_1kb", |b| {
        let mut codec = make_codec();
        let frame = Frame::Deliver {
            message_id: MessageId(42),
            offset: TopicOffset(100),
            timestamp_millis: 1_700_000_000_000,
            ttl_ms: 0,
            topic: TopicName("events.user.created".into()),
            headers: vec![],
            payload: Bytes::from(vec![0xCDu8; 1024]),
        };
        b.iter(|| {
            let mut buf = BytesMut::with_capacity(1200);
            codec.encode(frame.clone(), &mut buf).unwrap();
            codec.decode(&mut buf).unwrap().unwrap();
        });
    });

    group.finish();
}

fn codec_headers(c: &mut Criterion) {
    let mut group = c.benchmark_group("codec_headers");

    // Helper to make N headers with ~50 byte keys and ~100 byte values
    let make_headers = |n: usize| -> Vec<Header> {
        (0..n)
            .map(|i| Header {
                key: format!("x-custom-header-{:03}", i).into_bytes(),
                value: format!("value-{:0>96}", i).into_bytes(),
            })
            .collect()
    };

    for header_count in [0, 1, 5, 10] {
        let headers = make_headers(header_count);
        let label = format!("publish_1kb_{}h", header_count);

        // Encode with headers enabled
        group.bench_function(format!("encode_{}", label), |b| {
            let mut codec = make_codec();
            codec.set_headers_enabled(true);
            let frame = Frame::Publish {
                topic: TopicName("events".into()),
                headers: headers.clone(),
                ttl_ms: 0,
                payload: Bytes::from(vec![0xABu8; 1024]),
            };
            b.iter(|| {
                let mut buf = BytesMut::with_capacity(2048);
                codec.encode(frame.clone(), &mut buf).unwrap();
            });
        });

        // Decode with headers enabled
        group.bench_function(format!("decode_{}", label), |b| {
            let mut enc_codec = make_codec();
            enc_codec.set_headers_enabled(true);
            let frame = Frame::Publish {
                topic: TopicName("events".into()),
                headers: headers.clone(),
                ttl_ms: 0,
                payload: Bytes::from(vec![0xABu8; 1024]),
            };
            let mut encoded = BytesMut::new();
            enc_codec.encode(frame, &mut encoded).unwrap();
            let encoded_bytes = encoded.to_vec();

            let mut codec = make_codec();
            codec.set_headers_enabled(true);
            b.iter(|| {
                let mut buf = BytesMut::from(&encoded_bytes[..]);
                codec.decode(&mut buf).unwrap().unwrap();
            });
        });

        // Roundtrip with headers
        group.bench_function(format!("roundtrip_{}", label), |b| {
            let mut codec = make_codec();
            codec.set_headers_enabled(true);
            let frame = Frame::Publish {
                topic: TopicName("events".into()),
                headers: headers.clone(),
                ttl_ms: 0,
                payload: Bytes::from(vec![0xABu8; 1024]),
            };
            b.iter(|| {
                let mut buf = BytesMut::with_capacity(2048);
                codec.encode(frame.clone(), &mut buf).unwrap();
                codec.decode(&mut buf).unwrap().unwrap();
            });
        });
    }

    group.finish();
}

criterion_group!(
    benches,
    codec_encode,
    codec_decode,
    codec_roundtrip,
    codec_headers
);
criterion_main!(benches);
