path: root/examples/codec.rs

use zr_protocol::codec::decode::Decode as DecodeTrait;
use zr_protocol::codec::encode::Encode as EncodeTrait;
use zr_protocol::context::Context;

// ── Test Codec derive (combines Encode + Decode) ───────────────────────

#[derive(zr_protocol::Codec)]
#[context(Ctx)]
pub struct CodecStruct {
    id: u8,
    #[codec(count(u8))]
    items: Vec<u16>,
}

// ── State type for context-aware enums ───────────────────────────────

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum State {
    Handshake,
    Status,
    Play,
}

pub struct Ctx {
    pub state: State,
}

// ── Basic struct (no helpers) ────────────────────────────────────────

#[derive(zr_protocol::Encode, zr_protocol::Decode)]
pub struct BasicStruct {
    a: u8,
    b: u16,
    c: u32,
}

// ── Skip ─────────────────────────────────────────────────────────────

#[derive(zr_protocol::Encode)]
pub struct WithSkip {
    id: u8,
    #[codec(skip)]
    _cache: Vec<u8>,
    name: u16,
}

// ── Count prefix ─────────────────────────────────────────────────────

#[derive(zr_protocol::Encode, zr_protocol::Decode)]
pub struct WithCount {
    #[codec(count(u16))]
    items: Vec<u32>,
}

// ── Condition on struct field ────────────────────────────────────────

#[derive(zr_protocol::Encode)]
pub struct WithCondition {
    version: u8,
    #[codec(if(self.version >= 2))]
    extra: Option<u16>,
}

// ── Combined struct helpers ──────────────────────────────────────────

#[derive(zr_protocol::Encode)]
pub struct Combined {
    version: u8,
    #[codec(count(u16))]
    items: Vec<u32>,
    #[codec(if(self.items.is_empty()))]
    fallback: Option<u8>,
}

// ── Enum with auto IDs (basic) ──────────────────────────────────────

#[derive(zr_protocol::Encode, zr_protocol::Decode)]
pub enum BasicEnum {
    A(u8),
    B { x: u16, y: u16 },
    C,
}

// ── Enum with custom IDs ────────────────────────────────────────────

#[derive(zr_protocol::Encode, zr_protocol::Decode)]
pub enum CustomIdEnum {
    #[id(0x10)]
    Alpha(u8),
    #[id(0x20)]
    Beta { val: u16 },
    #[id(0x30)]
    Gamma,
}

// ── Context-aware enum (same ID, different states) ──────────────────

#[derive(zr_protocol::Encode, zr_protocol::Decode)]
#[context(Ctx)]
pub enum Packet {
    #[id(0x00)]
    #[codec(if(ctx.data.state == State::Handshake))]
    Handshake(u16),

    #[id(0x00)]
    #[codec(if(ctx.data.state == State::Status))]
    Status(u32),

    #[id(0x01)]
    #[codec(if(ctx.data.state == State::Play))]
    PlayData(u8, u8),

    #[id(0x02)]
    Ping,
}

// ── Context enum with skip ──────────────────────────────────────────

#[derive(zr_protocol::Encode, zr_protocol::Decode)]
#[context(Ctx)]
pub enum WithSkipVariant {
    #[id(0x00)]
    Visible(u8),
    #[codec(skip)]
    Internal(u16),
    #[id(0x01)]
    AlsoVisible,
}

// ── Test runner ──────────────────────────────────────────────────────

fn enc<T: EncodeTrait<Ctx>>(value: &T, state: State) -> Vec<u8> {
    let mut buf = Vec::new();
    let ctx = Context::new(Ctx { state });
    value.encode(&mut buf, &ctx).unwrap();
    buf
}

fn dec<T: DecodeTrait<Ctx>>(buf: &[u8], state: State) -> T {
    let ctx = Context::new(Ctx { state });
    T::decode(&mut &buf[..], &ctx).unwrap()
}

fn main() {
    // ── BasicStruct encode ──
    let basic = BasicStruct {
        a: 1,
        b: 0x203,
        c: 0x4050607,
    };
    let bytes = enc(&basic, State::Handshake);
    assert_eq!(bytes, vec![0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]);
    println!("BasicStruct encode: OK");

    // ── BasicStruct decode ──
    let decoded: BasicStruct = dec(&bytes, State::Handshake);
    assert_eq!(decoded.a, 1);
    assert_eq!(decoded.b, 0x203);
    assert_eq!(decoded.c, 0x4050607);
    println!("BasicStruct decode: OK");

    // ── WithSkip ──
    let skip = WithSkip {
        id: 0xAA,
        _cache: vec![1, 2, 3],
        name: 0xBBCC,
    };
    let bytes = enc(&skip, State::Handshake);
    assert_eq!(bytes, vec![0xAA, 0xBB, 0xCC]);
    println!("WithSkip:    OK");

    // ── WithCount encode ──
    let count_empty = WithCount { items: vec![] };
    let bytes = enc(&count_empty, State::Handshake);
    assert_eq!(bytes, vec![0x00, 0x00]);
    println!("WithCount (empty): OK");

    // ── WithCount decode ──
    let decoded: WithCount = dec(&bytes, State::Handshake);
    assert_eq!(decoded.items, Vec::<u32>::new());
    println!("WithCount (empty) decode: OK");

    // ── WithCount with items encode ──
    let count_some = WithCount {
        items: vec![0x0A, 0x0B],
    };
    let bytes = enc(&count_some, State::Handshake);
    assert_eq!(
        bytes,
        vec![0x00, 0x02, 0x00, 0x00, 0x00, 0x0A, 0x00, 0x00, 0x00, 0x0B]
    );
    println!("WithCount (2 items): OK");

    // ── WithCount with items decode ──
    let decoded: WithCount = dec(&bytes, State::Handshake);
    assert_eq!(decoded.items, vec![0x0A, 0x0B]);
    println!("WithCount (2 items) decode: OK");

    // ── WithCondition encode ──
    let cond_false = WithCondition {
        version: 1,
        extra: None,
    };
    let bytes = enc(&cond_false, State::Handshake);
    assert_eq!(bytes, vec![0x01]);
    println!("WithCondition (v1, None): OK");

    let cond_true = WithCondition {
        version: 2,
        extra: Some(0xDEAD),
    };
    let bytes = enc(&cond_true, State::Handshake);
    assert_eq!(bytes, vec![0x02, 0xDE, 0xAD]);
    println!("WithCondition (v2, Some): OK");

    // ── Combined encode ──
    let comb1 = Combined {
        version: 1,
        items: vec![42],
        fallback: Some(99),
    };
    let bytes = enc(&comb1, State::Handshake);
    assert_eq!(bytes, vec![0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x2A]);
    println!("Combined (items=[42]): OK");

    let comb2 = Combined {
        version: 1,
        items: vec![],
        fallback: Some(99),
    };
    let bytes = enc(&comb2, State::Handshake);
    assert_eq!(bytes, vec![0x01, 0x00, 0x00, 0x63]);
    println!("Combined (items=[], fallback): OK");

    // ── BasicEnum encode ──
    let e_a = BasicEnum::A(0xFF);
    let bytes = enc(&e_a, State::Handshake);
    assert_eq!(bytes, vec![0x00, 0xFF]);
    println!("Enum::A encode:     OK");

    let e_b = BasicEnum::B { x: 1, y: 2 };
    let bytes = enc(&e_b, State::Handshake);
    assert_eq!(bytes, vec![0x01, 0x00, 0x01, 0x00, 0x02]);
    println!("Enum::B encode:     OK");

    let e_c = BasicEnum::C;
    let bytes = enc(&e_c, State::Handshake);
    assert_eq!(bytes, vec![0x02]);
    println!("Enum::C encode:     OK");

    // ── BasicEnum decode ──
    let decoded: BasicEnum = dec(&[0x00, 0xFF], State::Handshake);
    assert!(matches!(decoded, BasicEnum::A(0xFF)));
    println!("Enum::A decode:     OK");

    let decoded: BasicEnum = dec(&[0x01, 0x00, 0x01, 0x00, 0x02], State::Handshake);
    assert!(matches!(decoded, BasicEnum::B { x: 1, y: 2 }));
    println!("Enum::B decode:     OK");

    let decoded: BasicEnum = dec(&[0x02], State::Handshake);
    assert!(matches!(decoded, BasicEnum::C));
    println!("Enum::C decode:     OK");

    // ── CustomIdEnum encode ──
    let bytes = enc(&CustomIdEnum::Alpha(42), State::Handshake);
    assert_eq!(bytes, vec![0x10, 42]);
    println!("CustomIdEnum::Alpha: OK");

    let bytes = enc(&CustomIdEnum::Beta { val: 1234 }, State::Handshake);
    assert_eq!(bytes, vec![0x20, 0x04, 0xD2]);
    println!("CustomIdEnum::Beta:  OK");

    let bytes = enc(&CustomIdEnum::Gamma, State::Handshake);
    assert_eq!(bytes, vec![0x30]);
    println!("CustomIdEnum::Gamma: OK");

    // ── CustomIdEnum decode ──
    let decoded: CustomIdEnum = dec(&[0x10, 42], State::Handshake);
    assert!(matches!(decoded, CustomIdEnum::Alpha(42)));
    println!("CustomIdEnum decode Alpha: OK");

    let decoded: CustomIdEnum = dec(&[0x20, 0x04, 0xD2], State::Handshake);
    assert!(matches!(decoded, CustomIdEnum::Beta { val: 1234 }));
    println!("CustomIdEnum decode Beta:  OK");

    let decoded: CustomIdEnum = dec(&[0x30], State::Handshake);
    assert!(matches!(decoded, CustomIdEnum::Gamma));
    println!("CustomIdEnum decode Gamma: OK");

    // ── Context-aware enum encode ──

    // Handshake state → ID 0x00 encodes Handshake
    let bytes = enc(&Packet::Handshake(42), State::Handshake);
    assert_eq!(bytes, vec![0x00, 0x00, 0x2A]);
    println!("Packet Handshake encode: OK ({:02x?}", bytes);

    // Status state → ID 0x00 encodes Status
    let bytes = enc(&Packet::Status(1000), State::Status);
    assert_eq!(bytes, vec![0x00, 0x00, 0x00, 0x03, 0xE8]);
    println!("Packet Status encode:    OK ({:02x?}", bytes);

    // Play state → ID 0x01 encodes PlayData
    let bytes = enc(&Packet::PlayData(1, 2), State::Play);
    assert_eq!(bytes, vec![0x01, 0x01, 0x02]);
    println!("Packet PlayData encode:  OK ({:02x?}", bytes);

    // Ping → always valid (no condition)
    let bytes = enc(&Packet::Ping, State::Handshake);
    assert_eq!(bytes, vec![0x02]);
    println!("Packet Ping encode:      OK");

    // ── Context-aware: wrong state → error ──
    let result = std::panic::catch_unwind(|| enc(&Packet::Handshake(42), State::Status));
    assert!(result.is_err(), "should panic: Handshake in Status state");
    println!("Packet Handshake in Status state: correctly errors");

    let result = std::panic::catch_unwind(|| enc(&Packet::Status(1000), State::Play));
    assert!(result.is_err(), "should panic: Status in Play state");
    println!("Packet Status in Play state:      correctly errors");

    // ── Context-aware enum decode ──

    // ID 0x00 in Handshake state → Handshake
    let decoded: Packet = dec(&[0x00, 0x00, 0x2A], State::Handshake);
    assert!(matches!(decoded, Packet::Handshake(42)));
    println!("Packet decode Handshake: OK");

    // ID 0x00 in Status state → Status
    let decoded: Packet = dec(&[0x00, 0x00, 0x00, 0x03, 0xE8], State::Status);
    assert!(matches!(decoded, Packet::Status(1000)));
    println!("Packet decode Status:    OK");

    // ID 0x01 in Play state → PlayData
    let decoded: Packet = dec(&[0x01, 0x01, 0x02], State::Play);
    assert!(matches!(decoded, Packet::PlayData(1, 2)));
    println!("Packet decode PlayData:  OK");

    // ID 0x02 → Ping (always valid)
    let decoded: Packet = dec(&[0x02], State::Handshake);
    assert!(matches!(decoded, Packet::Ping));
    println!("Packet decode Ping:      OK");

    // ── Context-aware decode: wrong state → error ──
    let result: Result<Packet, _> = {
        let ctx = Context::new(Ctx {
            state: State::Status,
        });
        let mut reader = &[0x00, 0x00, 0x2A][..];
        Packet::decode(&mut reader, &ctx)
    };
    assert!(
        result.is_err(),
        "should error: ID 0x00 in Status decodes as Status, not Handshake"
    );
    println!("Packet decode wrong context: correctly errors");

    // ── WithSkipVariant encode ──
    let bytes = enc(&WithSkipVariant::Visible(42), State::Handshake);
    assert_eq!(bytes, vec![0x00, 42]);
    println!("WithSkipVariant::Visible encode: OK");

    let bytes = enc(&WithSkipVariant::AlsoVisible, State::Handshake);
    assert_eq!(bytes, vec![0x01]);
    println!("WithSkipVariant::AlsoVisible encode: OK");

    // ── WithSkipVariant decode ──
    let decoded: WithSkipVariant = dec(&[0x00, 42], State::Handshake);
    assert!(matches!(decoded, WithSkipVariant::Visible(42)));
    println!("WithSkipVariant decode Visible: OK");

    let decoded: WithSkipVariant = dec(&[0x01], State::Handshake);
    assert!(matches!(decoded, WithSkipVariant::AlsoVisible));
    println!("WithSkipVariant decode AlsoVisible: OK");

    // Unknown discriminant → error
    let result: Result<WithSkipVariant, _> = {
        let ctx = Context::new(Ctx {
            state: State::Handshake,
        });
        let mut reader = &[0xFF][..];
        WithSkipVariant::decode(&mut reader, &ctx)
    };
    assert!(result.is_err(), "should error on unknown discriminant");
    println!("WithSkipVariant unknown discriminant: correctly errors");

    // ── Codec derive test ──
    let codec_val = CodecStruct {
        id: 0x42,
        items: vec![1, 2, 3],
    };
    let bytes = enc(&codec_val, State::Handshake);
    assert_eq!(bytes, vec![0x42, 0x03, 0x00, 0x01, 0x00, 0x02, 0x00, 0x03]);
    println!("CodecStruct encode: OK");

    let decoded: CodecStruct = dec(&bytes, State::Handshake);
    assert_eq!(decoded.id, 0x42);
    assert_eq!(decoded.items, vec![1, 2, 3]);
    println!("CodecStruct decode: OK");

    println!("\nAll tests passed!");
}