"""Tests for encoding.py boundary functions.

Verifies:
- pe-frame-redesign.AC2.2: pack_instruction / unpack_instruction round-trip
- Instruction encoding/decoding via hardware word format
- Flit 1 packing/unpacking for FrameDest
- Token packing/unpacking for T0 storage
"""

import pytest
from hypothesis import given, example, strategies as st

from cm_inst import (
    ArithOp, FrameDest, FrameOp, Instruction, LogicOp, MemOp,
    OutputStyle, Port, RoutingOp, TokenKind,
)
from encoding import (
    _decode_mode, _decode_opcode, _encode_mode, _encode_opcode,
    flit_count, pack_flit1, pack_instruction, pack_token, unpack_flit1,
    unpack_instruction, unpack_token,
)
from tokens import DyadToken, MonadToken, SMToken


# ============================================================================
# Mode Encoding/Decoding
# ============================================================================

class TestModeEncoding:
    """Mode field encoding matches design-notes table."""

    def test_encode_mode_inherit_single_no_const(self):
        """INHERIT + single output + no const = mode 0."""
        mode = _encode_mode(OutputStyle.INHERIT, False, 1)
        assert mode == 0b000

    def test_encode_mode_inherit_single_with_const(self):
        """INHERIT + single output + with const = mode 1."""
        mode = _encode_mode(OutputStyle.INHERIT, True, 1)
        assert mode == 0b001

    def test_encode_mode_inherit_fanout_no_const(self):
        """INHERIT + fan-out + no const = mode 2."""
        mode = _encode_mode(OutputStyle.INHERIT, False, 2)
        assert mode == 0b010

    def test_encode_mode_inherit_fanout_with_const(self):
        """INHERIT + fan-out + with const = mode 3."""
        mode = _encode_mode(OutputStyle.INHERIT, True, 2)
        assert mode == 0b011

    def test_encode_mode_change_tag_no_const(self):
        """CHANGE_TAG + no const = mode 4."""
        mode = _encode_mode(OutputStyle.CHANGE_TAG, False, 1)
        assert mode == 0b100

    def test_encode_mode_change_tag_with_const(self):
        """CHANGE_TAG + with const = mode 5."""
        mode = _encode_mode(OutputStyle.CHANGE_TAG, True, 1)
        assert mode == 0b101

    def test_encode_mode_sink_no_const(self):
        """SINK + no const = mode 6."""
        mode = _encode_mode(OutputStyle.SINK, False, 0)
        assert mode == 0b110

    def test_encode_mode_sink_with_const(self):
        """SINK + with const = mode 7."""
        mode = _encode_mode(OutputStyle.SINK, True, 0)
        assert mode == 0b111

    def test_encode_mode_invalid_inherit_dest_count(self):
        """INHERIT requires dest_count 1 or 2."""
        with pytest.raises(ValueError):
            _encode_mode(OutputStyle.INHERIT, False, 3)

    def test_encode_mode_invalid_output_style(self):
        """Invalid OutputStyle raises ValueError."""
        with pytest.raises(ValueError):
            _encode_mode(42, False, 1)  # type: ignore

    def test_encode_mode_change_tag_invalid_dest_count(self):
        """CHANGE_TAG requires dest_count == 1."""
        with pytest.raises(ValueError, match="CHANGE_TAG requires dest_count == 1"):
            _encode_mode(OutputStyle.CHANGE_TAG, False, 0)
        with pytest.raises(ValueError, match="CHANGE_TAG requires dest_count == 1"):
            _encode_mode(OutputStyle.CHANGE_TAG, False, 2)

    def test_encode_mode_sink_invalid_dest_count(self):
        """SINK requires dest_count == 0."""
        with pytest.raises(ValueError, match="SINK requires dest_count == 0"):
            _encode_mode(OutputStyle.SINK, False, 1)
        with pytest.raises(ValueError, match="SINK requires dest_count == 0"):
            _encode_mode(OutputStyle.SINK, False, 2)


class TestModeDecoding:
    """Mode field decoding extracts OutputStyle, has_const, dest_count."""

    def test_decode_mode_0(self):
        """Mode 0 decodes to INHERIT, no const, single output."""
        output, has_const, dest_count = _decode_mode(0b000)
        assert output == OutputStyle.INHERIT
        assert has_const is False
        assert dest_count == 1

    def test_decode_mode_1(self):
        """Mode 1 decodes to INHERIT, with const, single output."""
        output, has_const, dest_count = _decode_mode(0b001)
        assert output == OutputStyle.INHERIT
        assert has_const is True
        assert dest_count == 1

    def test_decode_mode_2(self):
        """Mode 2 decodes to INHERIT, no const, fan-out."""
        output, has_const, dest_count = _decode_mode(0b010)
        assert output == OutputStyle.INHERIT
        assert has_const is False
        assert dest_count == 2

    def test_decode_mode_3(self):
        """Mode 3 decodes to INHERIT, with const, fan-out."""
        output, has_const, dest_count = _decode_mode(0b011)
        assert output == OutputStyle.INHERIT
        assert has_const is True
        assert dest_count == 2

    def test_decode_mode_4(self):
        """Mode 4 decodes to CHANGE_TAG, no const."""
        output, has_const, dest_count = _decode_mode(0b100)
        assert output == OutputStyle.CHANGE_TAG
        assert has_const is False
        assert dest_count == 1  # nominal

    def test_decode_mode_5(self):
        """Mode 5 decodes to CHANGE_TAG, with const."""
        output, has_const, dest_count = _decode_mode(0b101)
        assert output == OutputStyle.CHANGE_TAG
        assert has_const is True
        assert dest_count == 1  # nominal

    def test_decode_mode_6(self):
        """Mode 6 decodes to SINK, no const."""
        output, has_const, dest_count = _decode_mode(0b110)
        assert output == OutputStyle.SINK
        assert has_const is False
        assert dest_count == 0

    def test_decode_mode_7(self):
        """Mode 7 decodes to SINK, with const."""
        output, has_const, dest_count = _decode_mode(0b111)
        assert output == OutputStyle.SINK
        assert has_const is True
        assert dest_count == 0


class TestModeRoundTrip:
    """Mode encoding round-trips correctly."""

    def test_mode_roundtrip_all_combinations(self):
        """All valid mode combinations round-trip."""
        for mode in range(8):
            output, has_const, dest_count = _decode_mode(mode)
            encoded = _encode_mode(output, has_const, dest_count)
            assert encoded == mode


# ============================================================================
# Opcode Encoding/Decoding
# ============================================================================

class TestOpcodeEncoding:
    """Opcode encoding and type_bit."""

    def test_encode_alu_opcode(self):
        """ALU opcodes use type_bit=0."""
        type_bit, opcode = _encode_opcode(ArithOp.ADD)
        assert type_bit == 0
        assert opcode == 0

    def test_encode_arith_ops(self):
        """ArithOp values encode correctly."""
        for op in [ArithOp.ADD, ArithOp.SUB, ArithOp.INC, ArithOp.DEC,
                   ArithOp.SHL, ArithOp.SHR, ArithOp.ASR]:
            type_bit, opcode = _encode_opcode(op)
            assert type_bit == 0
            assert opcode == int(op)

    def test_encode_logic_ops(self):
        """LogicOp values encode correctly."""
        for op in LogicOp:
            type_bit, opcode = _encode_opcode(op)
            assert type_bit == 0
            assert opcode == int(op)

    def test_encode_routing_ops(self):
        """RoutingOp values encode correctly."""
        for op in RoutingOp:
            type_bit, opcode = _encode_opcode(op)
            assert type_bit == 0
            assert opcode == int(op)

    def test_encode_memop(self):
        """MemOp opcodes use type_bit=1."""
        for op in MemOp:
            type_bit, opcode = _encode_opcode(op)
            assert type_bit == 1
            assert opcode == int(op)


class TestOpcodeDecoding:
    """Opcode decoding from type_bit + raw opcode."""

    def test_decode_alu_opcode(self):
        """ALU opcode decodes correctly."""
        op = _decode_opcode(0, 0)  # ArithOp.ADD
        assert op == ArithOp.ADD

    def test_decode_logic_opcode(self):
        """LogicOp decodes correctly."""
        op = _decode_opcode(0, 11)  # LogicOp.EQ
        assert op == LogicOp.EQ

    def test_decode_routing_opcode(self):
        """RoutingOp decodes correctly."""
        op = _decode_opcode(0, 16)  # RoutingOp.BREQ
        assert op == RoutingOp.BREQ

    def test_decode_memop(self):
        """MemOp decodes correctly."""
        op = _decode_opcode(1, 0)  # MemOp.READ
        assert op == MemOp.READ


class TestOpcodeRoundTrip:
    """Opcode encoding round-trips correctly."""

    def test_all_aluops_roundtrip(self):
        """All ALU ops round-trip."""
        for op in list(ArithOp) + list(LogicOp) + list(RoutingOp):
            type_bit, opcode = _encode_opcode(op)
            decoded = _decode_opcode(type_bit, opcode)
            assert decoded == op

    def test_all_memops_roundtrip(self):
        """All MemOps round-trip."""
        for op in MemOp:
            type_bit, opcode = _encode_opcode(op)
            decoded = _decode_opcode(type_bit, opcode)
            assert decoded == op


# ============================================================================
# Instruction Pack/Unpack
# ============================================================================

class TestInstructionPacking:
    """pack_instruction encodes to 16-bit word."""

    def test_pack_simple_add(self):
        """Simple ADD instruction packs correctly."""
        inst = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.INHERIT,
            has_const=False,
            dest_count=2,
            wide=False,
            fref=0,
        )
        word = pack_instruction(inst)
        # type_bit=0, opcode=0, mode=2, wide=0, fref=0
        # [0][00000][010][0][000000] = 0x0100
        assert word == 0x0100

    def test_pack_read_with_const(self):
        """READ instruction with const packs correctly."""
        inst = Instruction(
            opcode=MemOp.READ,
            output=OutputStyle.INHERIT,
            has_const=True,
            dest_count=1,
            wide=False,
            fref=10,
        )
        word = pack_instruction(inst)
        # type_bit=1, opcode=0, mode=1, wide=0, fref=10
        # [1][00000][001][0][001010]
        expected = (1 << 15) | (0 << 10) | (1 << 7) | (0 << 6) | 10
        assert word == expected

    def test_pack_wide_instruction(self):
        """Wide instructions pack with wide_bit=1."""
        inst = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.INHERIT,
            has_const=False,
            dest_count=2,
            wide=True,
            fref=0,
        )
        word = pack_instruction(inst)
        # wide_bit should be set
        assert (word >> 6) & 1 == 1

    def test_pack_max_fref(self):
        """Maximum fref (63) packs correctly."""
        inst = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.INHERIT,
            has_const=False,
            dest_count=1,
            wide=False,
            fref=63,
        )
        word = pack_instruction(inst)
        assert (word & 0x3F) == 63

    def test_pack_sink_output(self):
        """SINK output style encodes to mode 6 or 7."""
        inst = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.SINK,
            has_const=False,
            dest_count=0,
            wide=False,
            fref=0,
        )
        word = pack_instruction(inst)
        mode = (word >> 7) & 0x7
        assert mode == 0b110


class TestInstructionUnpacking:
    """unpack_instruction decodes 16-bit word to Instruction."""

    def test_unpack_simple_add(self):
        """Unpacking simple ADD gives correct Instruction."""
        word = 0x0100  # mode=2 (INHERIT, single, no const) -> actually dest_count=2
        inst = unpack_instruction(word)
        assert inst.opcode == ArithOp.ADD
        assert inst.output == OutputStyle.INHERIT
        assert inst.has_const is False
        assert inst.dest_count == 2
        assert inst.wide is False
        assert inst.fref == 0

    def test_unpack_read_with_const(self):
        """Unpacking READ with const gives correct Instruction."""
        word = (1 << 15) | (0 << 10) | (1 << 7) | (0 << 6) | 10
        inst = unpack_instruction(word)
        assert isinstance(inst.opcode, MemOp)
        assert inst.opcode == MemOp.READ
        assert inst.has_const is True
        assert inst.dest_count == 1
        assert inst.fref == 10

    def test_unpack_wide_bit(self):
        """Unpacking preserves wide flag."""
        word = 0x0080 | (1 << 6)  # Set wide_bit
        inst = unpack_instruction(word)
        assert inst.wide is True

    def test_unpack_change_tag_output(self):
        """Unpacking CHANGE_TAG mode gives correct output style."""
        word = (0 << 15) | (16 << 10) | (4 << 7)  # BREQ, mode 4
        inst = unpack_instruction(word)
        assert inst.output == OutputStyle.CHANGE_TAG


class TestInstructionRoundTrip:
    """pack_instruction and unpack_instruction round-trip."""

    @given(
        opcode=st.sampled_from(
            list(ArithOp) + list(LogicOp) + list(RoutingOp) + list(MemOp)
        ),
        output=st.sampled_from(list(OutputStyle)),
        has_const=st.booleans(),
        dest_count=st.integers(min_value=0, max_value=2),
        wide=st.booleans(),
        fref=st.integers(min_value=0, max_value=63),
    )
    # Boundary examples for instruction encoding
    @example(opcode=ArithOp.ADD, output=OutputStyle.INHERIT, has_const=False, dest_count=1, wide=False, fref=0)
    @example(opcode=ArithOp.ADD, output=OutputStyle.INHERIT, has_const=False, dest_count=2, wide=False, fref=63)
    @example(opcode=MemOp.READ, output=OutputStyle.INHERIT, has_const=True, dest_count=1, wide=True, fref=0)
    @example(opcode=RoutingOp.EXTRACT_TAG, output=OutputStyle.CHANGE_TAG, has_const=False, dest_count=1, wide=False, fref=31)
    @example(opcode=ArithOp.ADD, output=OutputStyle.SINK, has_const=False, dest_count=0, wide=False, fref=0)
    @example(opcode=ArithOp.ADD, output=OutputStyle.SINK, has_const=True, dest_count=0, wide=False, fref=63)
    def test_roundtrip_all_valid_combinations(self, opcode, output, has_const, dest_count, wide, fref):
        """All valid instruction combinations round-trip."""
        # Skip invalid combinations
        if output == OutputStyle.INHERIT and dest_count not in (1, 2):
            return
        if output == OutputStyle.SINK and dest_count != 0:
            return
        if output == OutputStyle.CHANGE_TAG:
            dest_count = 1  # CHANGE_TAG always decodes to dest_count=1

        inst = Instruction(
            opcode=opcode,
            output=output,
            has_const=has_const,
            dest_count=dest_count,
            wide=wide,
            fref=fref,
        )
        word = pack_instruction(inst)
        unpacked = unpack_instruction(word)

        assert unpacked.opcode == inst.opcode
        assert unpacked.output == inst.output
        assert unpacked.has_const == inst.has_const
        assert unpacked.dest_count == inst.dest_count
        assert unpacked.wide == inst.wide
        assert unpacked.fref == inst.fref


# ============================================================================
# Flit 1 Packing/Unpacking
# ============================================================================

class TestFlit1Packing:
    """pack_flit1 encodes FrameDest to 16-bit flit 1."""

    def test_pack_dyadic_dest(self):
        """DYADIC destination packs correctly."""
        dest = FrameDest(
            target_pe=1,
            offset=32,
            act_id=3,
            port=Port.L,
            token_kind=TokenKind.DYADIC,
        )
        flit1 = pack_flit1(dest)
        # Format: [00][port:1][PE:2][offset:8][act_id:3]
        # port=0, PE=01, offset=00100000, act_id=011
        assert (flit1 >> 14) == 0b00
        assert (flit1 >> 13) & 1 == 0  # Port.L = 0
        assert (flit1 >> 11) & 0x3 == 1  # PE = 1
        assert (flit1 >> 3) & 0xFF == 32  # offset
        assert flit1 & 0x7 == 3  # act_id

    def test_pack_monadic_dest(self):
        """MONADIC destination packs correctly."""
        dest = FrameDest(
            target_pe=2,
            offset=16,
            act_id=5,
            port=Port.R,  # ignored for MONADIC
            token_kind=TokenKind.MONADIC,
        )
        flit1 = pack_flit1(dest)
        # Format: [010][PE:2][offset:8][act_id:3]
        assert (flit1 >> 13) == 0b010
        assert (flit1 >> 11) & 0x3 == 2  # PE = 2
        assert (flit1 >> 3) & 0xFF == 16  # offset
        assert flit1 & 0x7 == 5  # act_id

    def test_pack_inline_dest(self):
        """INLINE destination packs correctly."""
        dest = FrameDest(
            target_pe=3,
            offset=64,
            act_id=0,  # ignored for INLINE
            port=Port.L,
            token_kind=TokenKind.INLINE,
        )
        flit1 = pack_flit1(dest)
        # Format: [011][PE:2][10][offset:7][spare:2]
        assert (flit1 >> 13) == 0b011
        assert (flit1 >> 11) & 0x3 == 3  # PE = 3
        assert (flit1 >> 9) & 0x3 == 0b10
        assert (flit1 >> 2) & 0x7F == 64  # offset


class TestFlit1Unpacking:
    """unpack_flit1 decodes 16-bit flit 1 to FrameDest."""

    def test_unpack_dyadic(self):
        """DYADIC flit unpacks correctly."""
        flit1 = (0b00 << 14) | (1 << 13) | (2 << 11) | (48 << 3) | 4
        dest = unpack_flit1(flit1)
        assert dest.token_kind == TokenKind.DYADIC
        assert dest.port == Port.R
        assert dest.target_pe == 2
        assert dest.offset == 48
        assert dest.act_id == 4

    def test_unpack_monadic(self):
        """MONADIC flit unpacks correctly."""
        flit1 = (0b010 << 13) | (1 << 11) | (20 << 3) | 2
        dest = unpack_flit1(flit1)
        assert dest.token_kind == TokenKind.MONADIC
        assert dest.port == Port.L
        assert dest.target_pe == 1
        assert dest.offset == 20
        assert dest.act_id == 2

    def test_unpack_inline(self):
        """INLINE flit unpacks correctly."""
        flit1 = (0b011 << 13) | (3 << 11) | (0b10 << 9) | (50 << 2)
        dest = unpack_flit1(flit1)
        assert dest.token_kind == TokenKind.INLINE
        assert dest.port == Port.L
        assert dest.target_pe == 3
        assert dest.offset == 50
        assert dest.act_id == 0


class TestFlit1RoundTrip:
    """pack_flit1 and unpack_flit1 round-trip."""

    @given(
        target_pe=st.integers(min_value=0, max_value=3),
        offset_dyadic=st.integers(min_value=0, max_value=255),
        offset_monadic=st.integers(min_value=0, max_value=255),
        offset_inline=st.integers(min_value=0, max_value=127),
        act_id=st.integers(min_value=0, max_value=7),
        port=st.sampled_from(list(Port)),
    )
    # Boundary examples for flit1 encoding
    @example(target_pe=0, offset_dyadic=0, offset_monadic=0, offset_inline=0, act_id=0, port=Port.L)
    @example(target_pe=3, offset_dyadic=255, offset_monadic=255, offset_inline=127, act_id=7, port=Port.R)
    @example(target_pe=1, offset_dyadic=128, offset_monadic=128, offset_inline=64, act_id=3, port=Port.L)
    @example(target_pe=2, offset_dyadic=1, offset_monadic=1, offset_inline=1, act_id=1, port=Port.R)
    def test_roundtrip_all_token_kinds(self, target_pe, offset_dyadic, offset_monadic, offset_inline, act_id, port):
        """All FrameDest combinations round-trip."""
        for token_kind, offset in [
            (TokenKind.DYADIC, offset_dyadic),
            (TokenKind.MONADIC, offset_monadic),
            (TokenKind.INLINE, offset_inline),
        ]:
            dest = FrameDest(
                target_pe=target_pe,
                offset=offset,
                act_id=act_id,
                port=port,
                token_kind=token_kind,
            )
            flit1 = pack_flit1(dest)
            unpacked = unpack_flit1(flit1)

            assert unpacked.target_pe == dest.target_pe
            assert unpacked.offset == dest.offset
            assert unpacked.token_kind == dest.token_kind
            # act_id should round-trip for all token kinds
            if token_kind != TokenKind.INLINE:
                assert unpacked.act_id == dest.act_id
            if token_kind == TokenKind.DYADIC:
                assert unpacked.port == dest.port
            else:
                assert unpacked.port == Port.L


# ============================================================================
# Token Packing/Unpacking
# ============================================================================

class TestTokenPacking:
    """pack_token encodes tokens to flit sequences."""

    def test_pack_dyadic_token(self):
        """DyadToken packs to 2 flits."""
        token = DyadToken(
            target=1,
            offset=32,
            act_id=3,
            data=0x1234,
            port=Port.L,
        )
        flits = pack_token(token)
        assert len(flits) == 2
        assert flits[1] == 0x1234

    def test_pack_monadic_token_normal(self):
        """Monadic (normal) token packs to 2 flits."""
        token = MonadToken(
            target=2,
            offset=16,
            act_id=5,
            data=0x5678,
            inline=False,
        )
        flits = pack_token(token)
        assert len(flits) == 2
        assert flits[1] == 0x5678

    def test_pack_monadic_token_inline(self):
        """Monadic (inline) token packs to 1 flit."""
        token = MonadToken(
            target=1,
            offset=64,
            act_id=0,
            data=0,  # ignored for inline
            inline=True,
        )
        flits = pack_token(token)
        assert len(flits) == 1

    def test_pack_smtoken(self):
        """SMToken packs to 2 flits."""
        token = SMToken(
            target=3,
            addr=100,
            op=MemOp.READ,
            flags=None,
            data=0xABCD,
            ret=None,
        )
        flits = pack_token(token)
        assert len(flits) == 2
        assert (flits[0] >> 15) & 1 == 1  # SM token marker
        assert flits[1] == 0xABCD

    def test_pack_smtoken_tier2_memop_raises(self):
        """SMToken with Tier 2 MemOp (value > 7) raises ValueError."""
        # Tier 2 MemOps: RD_DEC=8, CMP_SW=9, RAW_READ=10, SET_PAGE=11, WRITE_IMM=12
        # These cannot fit in 3 bits and pack_token should reject them
        token = SMToken(
            target=0,
            addr=50,
            op=MemOp.RD_DEC,  # value = 8, exceeds 3-bit limit
            flags=None,
            data=100,
            ret=None,
        )
        with pytest.raises(ValueError, match="exceeds 3-bit encoding limit"):
            pack_token(token)

    def test_pack_smtoken_all_tier2_memops_raise(self):
        """All Tier 2 MemOps raise ValueError on pack."""
        tier2_ops = [MemOp.RD_DEC, MemOp.CMP_SW, MemOp.RAW_READ, MemOp.SET_PAGE, MemOp.WRITE_IMM]
        for op in tier2_ops:
            if int(op) > 7:  # Only test if it's actually > 7
                token = SMToken(
                    target=0,
                    addr=0,
                    op=op,
                    flags=None,
                    data=0,
                    ret=None,
                )
                with pytest.raises(ValueError, match="not yet supported"):
                    pack_token(token)


class TestTokenUnpacking:
    """unpack_token decodes flit sequences to tokens."""

    def test_unpack_dyadic_token(self):
        """DyadToken unpacks from 2 flits."""
        token = DyadToken(
            target=1,
            offset=32,
            act_id=3,
            data=0x1234,
            port=Port.L,
        )
        flits = pack_token(token)
        unpacked = unpack_token(flits)

        assert isinstance(unpacked, DyadToken)
        assert unpacked.target == 1
        assert unpacked.offset == 32
        assert unpacked.act_id == 3
        assert unpacked.data == 0x1234
        assert unpacked.port == Port.L

    def test_unpack_monadic_normal(self):
        """Monadic (normal) token unpacks correctly."""
        token = MonadToken(
            target=2,
            offset=16,
            act_id=5,
            data=0x5678,
            inline=False,
        )
        flits = pack_token(token)
        unpacked = unpack_token(flits)

        assert isinstance(unpacked, MonadToken)
        assert unpacked.target == 2
        assert unpacked.offset == 16
        assert unpacked.act_id == 5
        assert unpacked.data == 0x5678
        assert unpacked.inline is False

    def test_unpack_monadic_inline(self):
        """Monadic (inline) token unpacks correctly."""
        token = MonadToken(
            target=1,
            offset=64,
            act_id=0,
            data=0,
            inline=True,
        )
        flits = pack_token(token)
        unpacked = unpack_token(flits)

        assert isinstance(unpacked, MonadToken)
        assert unpacked.inline is True
        assert unpacked.target == 1

    def test_unpack_smtoken(self):
        """SMToken unpacks correctly."""
        token = SMToken(
            target=3,
            addr=100,
            op=MemOp.READ,
            flags=None,
            data=0xABCD,
            ret=None,
        )
        flits = pack_token(token)
        unpacked = unpack_token(flits)

        assert isinstance(unpacked, SMToken)
        assert unpacked.target == 3
        assert unpacked.addr == 100
        assert unpacked.op == MemOp.READ
        assert unpacked.data == 0xABCD
        assert unpacked.ret is None  # SMToken.ret not preserved through pack/unpack


class TestTokenRoundTrip:
    """pack_token and unpack_token round-trip."""

    def test_dyadic_roundtrip(self):
        """DyadToken round-trips."""
        token = DyadToken(
            target=1,
            offset=32,
            act_id=3,
            data=0x1234,
            port=Port.R,
        )
        flits = pack_token(token)
        unpacked = unpack_token(flits)

        assert isinstance(unpacked, DyadToken)
        assert unpacked.target == token.target
        assert unpacked.offset == token.offset
        assert unpacked.act_id == token.act_id
        assert unpacked.data == token.data
        assert unpacked.port == token.port

    def test_monadic_roundtrip(self):
        """MonadToken round-trips."""
        token = MonadToken(
            target=2,
            offset=16,
            act_id=5,
            data=0x5678,
            inline=False,
        )
        flits = pack_token(token)
        unpacked = unpack_token(flits)

        assert isinstance(unpacked, MonadToken)
        assert unpacked.target == token.target
        assert unpacked.offset == token.offset
        assert unpacked.act_id == token.act_id
        assert unpacked.data == token.data
        assert unpacked.inline == token.inline

    def test_smtoken_roundtrip_except_ret(self):
        """SMToken round-trips (except ret field)."""
        token = SMToken(
            target=3,
            addr=100,
            op=MemOp.READ,
            flags=None,
            data=0xABCD,
            ret=None,  # ret is not preserved
        )
        flits = pack_token(token)
        unpacked = unpack_token(flits)

        assert isinstance(unpacked, SMToken)
        assert unpacked.target == token.target
        assert unpacked.addr == token.addr
        assert unpacked.op == token.op
        assert unpacked.data == token.data
        assert unpacked.ret is None


# ============================================================================
# Flit Count
# ============================================================================

class TestFlitCount:
    """flit_count determines packet size from flit 1."""

    def test_flit_count_dyadic(self):
        """Dyadic tokens are 2 flits."""
        dest = FrameDest(
            target_pe=1,
            offset=32,
            act_id=3,
            port=Port.L,
            token_kind=TokenKind.DYADIC,
        )
        flit1 = pack_flit1(dest)
        assert flit_count(flit1) == 2

    def test_flit_count_monadic_normal(self):
        """Monadic (normal) tokens are 2 flits."""
        dest = FrameDest(
            target_pe=1,
            offset=32,
            act_id=3,
            port=Port.L,
            token_kind=TokenKind.MONADIC,
        )
        flit1 = pack_flit1(dest)
        assert flit_count(flit1) == 2

    def test_flit_count_monadic_inline(self):
        """Monadic (inline) tokens are 1 flit."""
        dest = FrameDest(
            target_pe=1,
            offset=32,
            act_id=0,
            port=Port.L,
            token_kind=TokenKind.INLINE,
        )
        flit1 = pack_flit1(dest)
        assert flit_count(flit1) == 1

    def test_flit_count_smtoken(self):
        """SM tokens are 2 flits."""
        # SM token: [1][SM_id:2][op:3][addr:10]
        flit1 = (1 << 15) | (0 << 13) | (0 << 10) | 100
        assert flit_count(flit1) == 2