use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};
use zr_protocol::codec::{decode::Decode, encode::Encode};
use zr_protocol::context::Context;
use zr_protocol::types::prefix::{CountPrefix, LenPrefixed};
use zr_protocol_macros::Codec;

const DEFAULT_CTX: Context<()> = Context { data: () };

#[derive(Codec, Clone, Debug, PartialEq)]
struct BenchStruct {
    a: u8,
    b: u16,
    c: u32,
    d: u64,
    e: i8,
    f: i16,
    g: i32,
    h: i64,
    i: f32,
    j: f64,
    k: bool,
}

#[derive(Codec, Clone, Debug, PartialEq)]
struct BenchStructWithCollections {
    #[codec(count(u16))]
    items: Vec<u32>,
    #[codec(count(u16))]
    tags: Vec<String>,
    #[codec(len(u16))]
    payload: Vec<u8>,
}

#[derive(Codec, Clone, Debug, PartialEq)]
enum BenchEnum {
    A(u8),
    B { x: u16, y: u16 },
    C(u32, u32),
    D,
}

#[derive(Codec, Clone, Debug, PartialEq)]
struct ComplexStruct {
    id: u64,
    #[codec(count(u16))]
    values: Vec<i32>,
    #[codec(len(u16))]
    data: Vec<u8>,
    #[codec(count(u16))]
    names: Vec<String>,
    #[codec(count(u8))]
    flags: Vec<bool>,
    metadata: HashMap<String, u32>,
    tags: HashSet<String>,
}

use std::collections::{HashMap, HashSet};

fn encode_to_vec<T: Encode<()> + ?Sized>(value: &T) -> Vec<u8> {
    let mut buf = Vec::with_capacity(1024);
    value.encode(&mut buf, &DEFAULT_CTX).unwrap();
    buf
}

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

    macro_rules! bench_encode {
        ($name:expr, $value:expr) => {
            group.bench_function($name, |b| {
                b.iter(|| {
                    let mut buf = Vec::with_capacity(64);
                    black_box($value).encode(&mut buf, &DEFAULT_CTX).unwrap();
                    black_box(buf)
                })
            });
        };
    }

    bench_encode!("u8", &0xFFu8);
    bench_encode!("u16", &0xFFFFu16);
    bench_encode!("u32", &0xFFFFFFFFu32);
    bench_encode!("u64", &0xFFFFFFFFFFFFFFFFu64);
    bench_encode!("u128", &u128::MAX);
    bench_encode!("usize", &usize::MAX);
    bench_encode!("i8", &-128i8);
    bench_encode!("i16", &-32768i16);
    bench_encode!("i32", &-2147483648i32);
    bench_encode!("i64", &-9223372036854775808i64);
    bench_encode!("i128", &i128::MIN);
    bench_encode!("isize", &isize::MIN);
    bench_encode!("f32", &std::f32::consts::PI);
    bench_encode!("f64", &std::f64::consts::PI);
    bench_encode!("bool_true", &true);
    bench_encode!("bool_false", &false);
}

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

    macro_rules! bench_decode {
        ($name:expr, $bytes:expr, $ty:ty) => {
            group.bench_function($name, |b| {
                b.iter(|| {
                    let mut reader = &$bytes[..];
                    black_box(<$ty as Decode<()>>::decode(&mut reader, &DEFAULT_CTX).unwrap())
                })
            });
        };
    }

    bench_decode!("u8", encode_to_vec(&0xFFu8), u8);
    bench_decode!("u16", encode_to_vec(&0xFFFFu16), u16);
    bench_decode!("u32", encode_to_vec(&0xFFFFFFFFu32), u32);
    bench_decode!("u64", encode_to_vec(&0xFFFFFFFFFFFFFFFFu64), u64);
    bench_decode!("u128", encode_to_vec(&u128::MAX), u128);
    bench_decode!("usize", encode_to_vec(&usize::MAX), usize);
    bench_decode!("i8", encode_to_vec(&-128i8), i8);
    bench_decode!("i16", encode_to_vec(&-32768i16), i16);
    bench_decode!("i32", encode_to_vec(&-2147483648i32), i32);
    bench_decode!("i64", encode_to_vec(&-9223372036854775808i64), i64);
    bench_decode!("i128", encode_to_vec(&i128::MIN), i128);
    bench_decode!("isize", encode_to_vec(&isize::MIN), isize);
    bench_decode!("f32", encode_to_vec(&std::f32::consts::PI), f32);
    bench_decode!("f64", encode_to_vec(&std::f64::consts::PI), f64);
    bench_decode!("bool_true", encode_to_vec(&true), bool);
    bench_decode!("bool_false", encode_to_vec(&false), bool);
}

fn bench_struct_encode(c: &mut Criterion) {
    let mut group = c.benchmark_group("struct_encode");
    let val = BenchStruct {
        a: 0xFF,
        b: 0xFFFF,
        c: 0xFFFFFFFF,
        d: 0xFFFFFFFFFFFFFFFF,
        e: -128,
        f: -32768,
        g: -2147483648,
        h: -9223372036854775808,
        i: std::f32::consts::PI,
        j: std::f64::consts::PI,
        k: true,
    };

    group.bench_function("encode", |b| {
        b.iter(|| {
            let mut buf = Vec::with_capacity(128);
            black_box(&val).encode(&mut buf, &DEFAULT_CTX).unwrap();
            black_box(buf)
        })
    });
}

fn bench_struct_decode(c: &mut Criterion) {
    let mut group = c.benchmark_group("struct_decode");
    let val = BenchStruct {
        a: 0xFF,
        b: 0xFFFF,
        c: 0xFFFFFFFF,
        d: 0xFFFFFFFFFFFFFFFF,
        e: -128,
        f: -32768,
        g: -2147483648,
        h: -9223372036854775808,
        i: std::f32::consts::PI,
        j: std::f64::consts::PI,
        k: true,
    };
    let bytes = encode_to_vec(&val);

    group.bench_function("decode", |b| {
        b.iter(|| {
            let mut reader = &bytes[..];
            black_box(BenchStruct::decode(&mut reader, &DEFAULT_CTX).unwrap())
        })
    });
}

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

    for size in [0, 10, 100, 1000, 10000].iter() {
        let items: Vec<u32> = (0..*size).collect();
        let tags: Vec<String> = (0..*size).map(|i| format!("tag{}", i)).collect();
        let payload: Vec<u8> = (0..*size).map(|i| (i % 256) as u8).collect();

        let val = BenchStructWithCollections {
            items: items.clone(),
            tags: tags.clone(),
            payload: payload.clone(),
        };

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(
            BenchmarkId::new("Vec_u32_count", size),
            &items,
            |b, items| {
                b.iter(|| {
                    let mut buf = Vec::with_capacity(1024);
                    CountPrefix::<u32, u16, _>::new(items.clone())
                        .encode(&mut buf, &DEFAULT_CTX)
                        .unwrap();
                    black_box(buf)
                })
            },
        );

        group.bench_with_input(
            BenchmarkId::new("Vec_String_count", size),
            &tags,
            |b, tags| {
                b.iter(|| {
                    let mut buf = Vec::with_capacity(1024);
                    CountPrefix::<String, u16, _>::new(tags.clone())
                        .encode(&mut buf, &DEFAULT_CTX)
                        .unwrap();
                    black_box(buf)
                })
            },
        );

        group.bench_with_input(
            BenchmarkId::new("Vec_u8_len", size),
            &payload,
            |b, payload| {
                b.iter(|| {
                    let mut buf = Vec::with_capacity(1024);
                    LenPrefixed::<u16, _>::new(payload.clone())
                        .encode(&mut buf, &DEFAULT_CTX)
                        .unwrap();
                    black_box(buf)
                })
            },
        );

        group.bench_with_input(BenchmarkId::new("full_struct", size), &val, |b, val| {
            b.iter(|| {
                let mut buf = Vec::with_capacity(4096);
                black_box(val).encode(&mut buf, &DEFAULT_CTX).unwrap();
                black_box(buf)
            })
        });
    }
}

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

    for size in [0, 10, 100, 1000, 10000].iter() {
        let items: Vec<u32> = (0..*size).collect();
        let tags: Vec<String> = (0..*size).map(|i| format!("tag{}", i)).collect();
        let payload: Vec<u8> = (0..*size).map(|i| (i % 256) as u8).collect();

        let val = BenchStructWithCollections {
            items: items.clone(),
            tags: tags.clone(),
            payload: payload.clone(),
        };
        let bytes = encode_to_vec(&val);

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(
            BenchmarkId::new("Vec_u32_count", size),
            &bytes,
            |b, bytes| {
                b.iter(|| {
                    let mut reader = &bytes[..];
                    black_box(
                        CountPrefix::<u32, u16, Vec<u32>>::decode(&mut reader, &DEFAULT_CTX)
                            .unwrap(),
                    )
                })
            },
        );

        group.bench_with_input(
            BenchmarkId::new("Vec_String_count", size),
            &bytes,
            |b, bytes| {
                b.iter(|| {
                    let mut reader = &bytes[..];
                    black_box(
                        CountPrefix::<String, u16, Vec<String>>::decode(&mut reader, &DEFAULT_CTX)
                            .unwrap(),
                    )
                })
            },
        );

        group.bench_with_input(BenchmarkId::new("Vec_u8_len", size), &bytes, |b, bytes| {
            b.iter(|| {
                let mut reader = &bytes[..];
                black_box(LenPrefixed::<u16, Vec<u8>>::decode(&mut reader, &DEFAULT_CTX).unwrap())
            })
        });

        group.bench_with_input(BenchmarkId::new("full_struct", size), &bytes, |b, bytes| {
            b.iter(|| {
                let mut reader = &bytes[..];
                black_box(BenchStructWithCollections::decode(&mut reader, &DEFAULT_CTX).unwrap())
            })
        });
    }
}

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

    for size in [0, 10, 100, 1000, 10000].iter() {
        let s = "x".repeat(*size);
        let s_clone = s.clone();

        group.throughput(Throughput::Bytes(*size as u64));

        group.bench_with_input(BenchmarkId::new("String_raw", size), &s_clone, |b, s| {
            b.iter(|| {
                let mut buf = Vec::with_capacity(*size + 8);
                black_box(s).encode(&mut buf, &DEFAULT_CTX).unwrap();
                black_box(buf)
            })
        });

        group.bench_with_input(BenchmarkId::new("str_raw", size), &s, |b, s| {
            b.iter(|| {
                let mut buf = Vec::with_capacity(*size + 8);
                black_box(s.as_str())
                    .encode(&mut buf, &DEFAULT_CTX)
                    .unwrap();
                black_box(buf)
            })
        });

        group.bench_with_input(
            BenchmarkId::new("LenPrefixed_String", size),
            &s_clone,
            |b, s| {
                b.iter(|| {
                    let mut buf = Vec::with_capacity(*size + 8);
                    LenPrefixed::<u16, _>::new(s.clone())
                        .encode(&mut buf, &DEFAULT_CTX)
                        .unwrap();
                    black_box(buf)
                })
            },
        );
    }
}

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

    for size in [0, 10, 100, 1000, 10000].iter() {
        let s = "x".repeat(*size);

        let raw_bytes = encode_to_vec(&s);
        let len_prefixed_bytes = encode_to_vec(&LenPrefixed::<u16, _>::new(s.clone()));

        group.throughput(Throughput::Bytes(*size as u64));

        group.bench_with_input(
            BenchmarkId::new("String_raw", size),
            &raw_bytes,
            |b, bytes| {
                b.iter(|| {
                    let mut reader = &bytes[..];
                    black_box(String::decode(&mut reader, &DEFAULT_CTX).unwrap())
                })
            },
        );

        group.bench_with_input(
            BenchmarkId::new("LenPrefixed_String", size),
            &len_prefixed_bytes,
            |b, bytes| {
                b.iter(|| {
                    let mut reader = &bytes[..];
                    black_box(
                        LenPrefixed::<u16, String>::decode(&mut reader, &DEFAULT_CTX).unwrap(),
                    )
                })
            },
        );
    }
}

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

    let variants = [
        ("A", BenchEnum::A(0xFF)),
        (
            "B",
            BenchEnum::B {
                x: 0xFFFF,
                y: 0xFFFF,
            },
        ),
        ("C", BenchEnum::C(0xFFFFFFFF, 0xFFFFFFFF)),
        ("D", BenchEnum::D),
    ];

    for (name, val) in variants {
        group.bench_function(name, |b| {
            b.iter(|| {
                let mut buf = Vec::with_capacity(32);
                black_box(&val).encode(&mut buf, &DEFAULT_CTX).unwrap();
                black_box(buf)
            })
        });
    }
}

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

    let variants = [
        ("A", encode_to_vec(&BenchEnum::A(0xFF))),
        (
            "B",
            encode_to_vec(&BenchEnum::B {
                x: 0xFFFF,
                y: 0xFFFF,
            }),
        ),
        ("C", encode_to_vec(&BenchEnum::C(0xFFFFFFFF, 0xFFFFFFFF))),
        ("D", encode_to_vec(&BenchEnum::D)),
    ];

    for (name, bytes) in variants {
        group.bench_function(name, |b| {
            b.iter(|| {
                let mut reader = &bytes[..];
                black_box(BenchEnum::decode(&mut reader, &DEFAULT_CTX).unwrap())
            })
        });
    }
}

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

    let struct_val = BenchStruct {
        a: 0xFF,
        b: 0xFFFF,
        c: 0xFFFFFFFF,
        d: 0xFFFFFFFFFFFFFFFF,
        e: -128,
        f: -32768,
        g: -2147483648,
        h: -9223372036854775808,
        i: std::f32::consts::PI,
        j: std::f64::consts::PI,
        k: true,
    };

    group.bench_function("struct", |b| {
        b.iter(|| {
            let mut buf = Vec::with_capacity(128);
            struct_val.encode(&mut buf, &DEFAULT_CTX).unwrap();
            let mut reader = &buf[..];
            black_box(BenchStruct::decode(&mut reader, &DEFAULT_CTX).unwrap())
        })
    });

    for size in [0, 10, 100, 1000].iter() {
        let items: Vec<u32> = (0..*size).collect();
        let val = CountPrefix::<u32, u16, _>::new(items);

        group.throughput(Throughput::Elements(*size as u64));
        group.bench_with_input(
            BenchmarkId::new("CountPrefix_Vec_u32", size),
            &val,
            |b, val| {
                b.iter(|| {
                    let mut buf = Vec::with_capacity(4096);
                    val.encode(&mut buf, &DEFAULT_CTX).unwrap();
                    let mut reader = &buf[..];
                    black_box(
                        CountPrefix::<u32, u16, Vec<u32>>::decode(&mut reader, &DEFAULT_CTX)
                            .unwrap(),
                    )
                })
            },
        );
    }
}

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

    for size in [0, 10, 100, 1000].iter() {
        let mut values = Vec::with_capacity(*size);
        let mut data = Vec::with_capacity(*size * 4);
        let mut names = Vec::with_capacity(*size);
        let mut flags = Vec::with_capacity(*size);
        let mut metadata = HashMap::new();
        let mut tags = HashSet::new();

        for i in 0..*size {
            values.push(i as i32);
            data.extend_from_slice(&(i as u32).to_be_bytes());
            names.push(format!("name_{}", i));
            flags.push(i % 2 == 0);
            metadata.insert(format!("key_{}", i), i as u32);
            tags.insert(format!("tag_{}", i));
        }

        let val = ComplexStruct {
            id: 0xDEADBEEF,
            values,
            data,
            names,
            flags,
            metadata,
            tags,
        };

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(BenchmarkId::new("ComplexStruct", size), &val, |b, val| {
            b.iter(|| {
                let mut buf = Vec::with_capacity(16384);
                black_box(val).encode(&mut buf, &DEFAULT_CTX).unwrap();
                black_box(buf)
            })
        });
    }
}

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

    for size in [0, 10, 100, 1000].iter() {
        let mut values = Vec::with_capacity(*size);
        let mut data = Vec::with_capacity(*size * 4);
        let mut names = Vec::with_capacity(*size);
        let mut flags = Vec::with_capacity(*size);
        let mut metadata = HashMap::new();
        let mut tags = HashSet::new();

        for i in 0..*size {
            values.push(i as i32);
            data.extend_from_slice(&(i as u32).to_be_bytes());
            names.push(format!("name_{}", i));
            flags.push(i % 2 == 0);
            metadata.insert(format!("key_{}", i), i as u32);
            tags.insert(format!("tag_{}", i));
        }

        let val = ComplexStruct {
            id: 0xDEADBEEF,
            values,
            data,
            names,
            flags,
            metadata,
            tags,
        };
        let bytes = encode_to_vec(&val);

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(
            BenchmarkId::new("ComplexStruct", size),
            &bytes,
            |b, bytes| {
                b.iter(|| {
                    let mut reader = &bytes[..];
                    black_box(ComplexStruct::decode(&mut reader, &DEFAULT_CTX).unwrap())
                })
            },
        );
    }
}

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

    for size in [10, 100, 1000, 10000].iter() {
        let data: Vec<u32> = (0..*size).collect();

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(BenchmarkId::new("slice_u32", size), &data, |b, data| {
            b.iter(|| {
                let mut buf = Vec::with_capacity((*size * 4 + 8) as usize);
                <_ as Encode<()>>::encode(data, &mut buf, &DEFAULT_CTX).unwrap();
                black_box(buf)
            })
        });
    }
}

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

    for size in [10, 100, 1000, 10000].iter() {
        let data: Vec<u32> = (0..*size).collect();
        let bytes = encode_to_vec(&data);

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(BenchmarkId::new("Vec_u32", size), &bytes, |b, bytes| {
            b.iter(|| {
                let mut reader = &bytes[..];
                black_box(<Vec<u32> as Decode<()>>::decode(&mut reader, &DEFAULT_CTX).unwrap())
            })
        });
    }
}

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

    for size in [0, 10, 100, 1000].iter() {
        let mut map = HashMap::with_capacity(*size);
        for i in 0..*size {
            map.insert(format!("key_{}", i), i as u32);
        }

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(
            BenchmarkId::new("HashMap_String_u32", size),
            &map,
            |b, map| {
                b.iter(|| {
                    let mut buf = Vec::with_capacity(4096);
                    black_box(map).encode(&mut buf, &DEFAULT_CTX).unwrap();
                    black_box(buf)
                })
            },
        );
    }
}

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

    for size in [0, 10, 100, 1000].iter() {
        let mut map = HashMap::with_capacity(*size);
        for i in 0..*size {
            map.insert(format!("key_{}", i), i as u32);
        }
        let bytes = encode_to_vec(&map);

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(
            BenchmarkId::new("HashMap_String_u32", size),
            &bytes,
            |b, bytes| {
                b.iter(|| {
                    let mut reader = &bytes[..];
                    black_box(
                        <HashMap<String, u32> as Decode<()>>::decode(&mut reader, &DEFAULT_CTX)
                            .unwrap(),
                    )
                })
            },
        );
    }
}

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

    for size in [0, 10, 100, 1000].iter() {
        let mut set = HashSet::with_capacity(*size);
        for i in 0..*size {
            set.insert(format!("item_{}", i));
        }

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(BenchmarkId::new("HashSet_String", size), &set, |b, set| {
            b.iter(|| {
                let mut buf = Vec::with_capacity(4096);
                black_box(set).encode(&mut buf, &DEFAULT_CTX).unwrap();
                black_box(buf)
            })
        });
    }
}

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

    for size in [0, 10, 100, 1000].iter() {
        let mut set = HashSet::with_capacity(*size);
        for i in 0..*size {
            set.insert(format!("item_{}", i));
        }
        let bytes = encode_to_vec(&set);

        group.throughput(Throughput::Elements(*size as u64));

        group.bench_with_input(
            BenchmarkId::new("HashSet_String", size),
            &bytes,
            |b, bytes| {
                b.iter(|| {
                    let mut reader = &bytes[..];
                    black_box(
                        <HashSet<String> as Decode<()>>::decode(&mut reader, &DEFAULT_CTX).unwrap(),
                    )
                })
            },
        );
    }
}

criterion_group!(
    benches,
    bench_primitives_encode,
    bench_primitives_decode,
    bench_struct_encode,
    bench_struct_decode,
    bench_collections_encode,
    bench_collections_decode,
    bench_string_encode,
    bench_string_decode,
    bench_enum_encode,
    bench_enum_decode,
    bench_roundtrip,
    bench_complex_struct_encode,
    bench_complex_struct_decode,
    bench_slice_encode,
    bench_slice_decode,
    bench_hashmap_encode,
    bench_hashmap_decode,
    bench_hashset_encode,
    bench_hashset_decode,
);
criterion_main!(benches);