"""
Tests for SM T0 raw int storage and EXEC flit-based parsing.

Verifies acceptance criteria:
- pe-frame-redesign.AC4.1: SM T0 stores list[int] (16-bit words), not Token objects
- pe-frame-redesign.AC4.2: EXEC reads consecutive ints, uses flit_count() for packet boundaries,
  reconstitutes tokens via unpack_token()
- pe-frame-redesign.AC4.3: pack_token() / unpack_token() round-trip for all token types
- pe-frame-redesign.AC4.4: Malformed flit sequence in T0 (invalid prefix bits) is handled
  gracefully without crash
"""

import pytest
import simpy

from cm_inst import MemOp, Port
from emu.sm import StructureMemory
from encoding import pack_token, unpack_token
from tokens import CMToken, DyadToken, MonadToken, SMToken


class TestAC4_1T0RawIntStorage:
    """AC4.1: SM T0 stores list[int] (16-bit words), not Token objects."""

    def test_t0_store_initialized_as_empty_int_list(self):
        """t0_store is initialized as empty list[int]."""
        env = simpy.Environment()
        sm = StructureMemory(env, 0, cell_count=512, tier_boundary=256)

        assert isinstance(sm.t0_store, list)
        assert len(sm.t0_store) == 0
        # Should be list of ints, not list of Tokens

    def test_t0_write_stores_int_values(self):
        """T0 WRITE stores int values."""
        env = simpy.Environment()
        sm = StructureMemory(env, 0, cell_count=512, tier_boundary=256)
        collector = simpy.Store(env)
        sm.route_table[0] = collector

        # Write int to T0
        write_token = SMToken(
            target=0, addr=256, op=MemOp.WRITE,
            flags=None, data=0x1234, ret=None
        )

        def do_write():
            yield sm.input_store.put(write_token)

        env.process(do_write())
        env.run()

        # Verify T0[0] (addr 256 = tier_boundary) now contains 0x1234
        assert len(sm.t0_store) >= 1
        assert sm.t0_store[0] == 0x1234


class TestAC4_2ExecFlitBasedParsing:
    """AC4.2: EXEC uses flit_count() for packet boundaries, unpack_token() to reconstitute."""

    def test_exec_reconstitutes_single_dyadic_token(self):
        """EXEC reads packed dyadic token flits, reconstitutes it."""
        env = simpy.Environment()
        sm = StructureMemory(env, 0, cell_count=512, tier_boundary=256)

        # Mock system and PE input store for token delivery
        from unittest.mock import Mock
        mock_pe_store = simpy.Store(env)
        mock_system = Mock()
        sm.system = mock_system

        # Set up system.send() to deliver to mock_pe_store
        def mock_send(token):
            yield env.timeout(1)
            yield mock_pe_store.put(token)

        mock_system.send = mock_send

        # Create a dyadic token and pack it
        orig_token = DyadToken(
            target=0, offset=10, act_id=2, data=0xABCD, port=Port.L
        )
        flits = pack_token(orig_token)

        # Pre-load T0 with packed flits
        sm.t0_store.extend(flits)

        # Trigger EXEC at T0 address 256 (tier_boundary)
        exec_token = SMToken(
            target=0, addr=256, op=MemOp.EXEC,
            flags=None, data=None, ret=None
        )

        def do_exec():
            yield sm.input_store.put(exec_token)

        env.process(do_exec())
        env.run()

        # Verify reconstituted token arrived at PE
        assert len(mock_pe_store.items) >= 1
        received = mock_pe_store.items[0]
        assert isinstance(received, DyadToken)
        assert received.target == 0
        assert received.offset == 10
        assert received.act_id == 2
        assert received.data == 0xABCD
        assert received.port == Port.L

    def test_exec_reconstitutes_multiple_consecutive_packets(self):
        """EXEC parses multiple consecutive token packets by flit_count boundary."""
        env = simpy.Environment()
        sm = StructureMemory(env, 0, cell_count=512, tier_boundary=256)

        # Mock system and PE input store
        from unittest.mock import Mock
        mock_pe_store = simpy.Store(env)
        mock_system = Mock()
        sm.system = mock_system

        def mock_send(token):
            yield env.timeout(1)
            yield mock_pe_store.put(token)

        mock_system.send = mock_send

        # Create two tokens
        token1 = DyadToken(target=0, offset=5, act_id=1, data=0x1111, port=Port.L)
        token2 = MonadToken(target=0, offset=15, act_id=3, data=0x2222, inline=False)

        flits1 = pack_token(token1)
        flits2 = pack_token(token2)

        # Pre-load both in sequence
        sm.t0_store.extend(flits1)
        sm.t0_store.extend(flits2)

        # EXEC from start of T0
        exec_token = SMToken(
            target=0, addr=256, op=MemOp.EXEC,
            flags=None, data=None, ret=None
        )

        def do_exec():
            yield sm.input_store.put(exec_token)

        env.process(do_exec())
        env.run()

        # Verify both tokens arrived
        assert len(mock_pe_store.items) >= 2
        received1 = mock_pe_store.items[0]
        received2 = mock_pe_store.items[1]

        assert isinstance(received1, DyadToken)
        assert received1.data == 0x1111

        assert isinstance(received2, MonadToken)
        assert received2.data == 0x2222

    def test_exec_processes_multiple_packets_until_end(self):
        """EXEC processes multiple consecutive packets until end of T0 store."""
        env = simpy.Environment()
        sm = StructureMemory(env, 0, cell_count=512, tier_boundary=256)

        # Mock system and PE input store
        from unittest.mock import Mock
        mock_pe_store = simpy.Store(env)
        mock_system = Mock()
        sm.system = mock_system

        def mock_send(token):
            yield env.timeout(1)
            yield mock_pe_store.put(token)

        mock_system.send = mock_send

        # Create three valid tokens
        token1 = DyadToken(target=0, offset=5, act_id=1, data=0x1111, port=Port.L)
        token2 = MonadToken(target=1, offset=10, act_id=2, data=0x2222, inline=False)
        token3 = DyadToken(target=2, offset=15, act_id=3, data=0x3333, port=Port.R)

        flits1 = pack_token(token1)
        flits2 = pack_token(token2)
        flits3 = pack_token(token3)

        # Pre-load all three tokens
        sm.t0_store.extend(flits1)
        sm.t0_store.extend(flits2)
        sm.t0_store.extend(flits3)

        # EXEC from start
        exec_token = SMToken(
            target=0, addr=256, op=MemOp.EXEC,
            flags=None, data=None, ret=None
        )

        def do_exec():
            yield sm.input_store.put(exec_token)

        env.process(do_exec())
        env.run()

        # Should have received all three tokens
        assert len(mock_pe_store.items) == 3
        assert mock_pe_store.items[0].data == 0x1111
        assert mock_pe_store.items[1].data == 0x2222
        assert mock_pe_store.items[2].data == 0x3333


class TestAC4_3RoundTripPackUnpack:
    """AC4.3: pack_token() / unpack_token() round-trip for all token types."""

    def test_dyadic_round_trip(self):
        """pack_token() -> unpack_token() preserves DyadToken fields."""
        token = DyadToken(
            target=2, offset=42, act_id=5, data=0x3456, port=Port.R
        )
        flits = pack_token(token)
        reconstituted = unpack_token(flits)

        assert isinstance(reconstituted, DyadToken)
        assert reconstituted.target == 2
        assert reconstituted.offset == 42
        assert reconstituted.act_id == 5
        assert reconstituted.data == 0x3456
        assert reconstituted.port == Port.R

    def test_monadic_normal_round_trip(self):
        """pack_token() -> unpack_token() preserves MonadToken (normal) fields."""
        token = MonadToken(
            target=1, offset=30, act_id=2, data=0x7890, inline=False
        )
        flits = pack_token(token)
        reconstituted = unpack_token(flits)

        assert isinstance(reconstituted, MonadToken)
        assert reconstituted.target == 1
        assert reconstituted.offset == 30
        assert reconstituted.act_id == 2
        assert reconstituted.data == 0x7890
        assert reconstituted.inline == False

    def test_monadic_inline_round_trip(self):
        """pack_token() -> unpack_token() preserves MonadToken (inline) fields."""
        token = MonadToken(
            target=0, offset=20, act_id=1, data=0, inline=True
        )
        flits = pack_token(token)
        assert len(flits) == 1  # inline has only 1 flit
        reconstituted = unpack_token(flits)

        assert isinstance(reconstituted, MonadToken)
        assert reconstituted.target == 0
        assert reconstituted.offset == 20
        # Note: act_id is not encoded in inline format (hardware constraint)
        assert reconstituted.act_id == 0
        assert reconstituted.inline == True

    def test_sm_token_round_trip(self):
        """pack_token() -> unpack_token() preserves SMToken fields."""
        token = SMToken(
            target=3, addr=512, op=MemOp.READ, flags=None, data=0x4567, ret=None
        )
        flits = pack_token(token)
        reconstituted = unpack_token(flits)

        assert isinstance(reconstituted, SMToken)
        assert reconstituted.target == 3
        assert reconstituted.addr == 512
        assert reconstituted.op == MemOp.READ
        assert reconstituted.data == 0x4567


class TestAC4_4GracefulErrorHandling:
    """AC4.4: Malformed flit sequences are handled gracefully without crash."""

    def test_malformed_flit_invalid_prefix(self):
        """EXEC stops gracefully when flit count exceeds available data (truncation detection)."""
        env = simpy.Environment()
        sm = StructureMemory(env, 0, cell_count=512, tier_boundary=256)

        # Mock system and PE input store
        from unittest.mock import Mock
        mock_pe_store = simpy.Store(env)
        mock_system = Mock()
        sm.system = mock_system

        def mock_send(token):
            yield env.timeout(1)
            yield mock_pe_store.put(token)

        mock_system.send = mock_send

        # Pre-load T0 with 0xFFFF (SM token prefix).
        # flit_count(0xFFFF) returns 2 (SM token header indicates 2 flits),
        # but only 1 flit is in store, so truncation check catches it.
        sm.t0_store.extend([0xFFFF])

        # EXEC from start
        exec_token = SMToken(
            target=0, addr=256, op=MemOp.EXEC,
            flags=None, data=None, ret=None
        )

        def do_exec():
            yield sm.input_store.put(exec_token)

        env.process(do_exec())
        # Should not crash
        env.run()

        # No tokens should have been injected (parse failed)
        assert len(mock_pe_store.items) == 0

    def test_truncated_packet_stops_gracefully(self):
        """EXEC stops gracefully when packet is truncated."""
        env = simpy.Environment()
        sm = StructureMemory(env, 0, cell_count=512, tier_boundary=256)

        # Mock system and PE input store
        from unittest.mock import Mock
        mock_pe_store = simpy.Store(env)
        mock_system = Mock()
        sm.system = mock_system

        def mock_send(token):
            yield env.timeout(1)
            yield mock_pe_store.put(token)

        mock_system.send = mock_send

        # Create a 2-flit dyadic token
        token = DyadToken(target=0, offset=5, act_id=1, data=0x1111, port=Port.L)
        flits = pack_token(token)
        assert len(flits) == 2

        # Pre-load with first flit only (truncated)
        sm.t0_store.append(flits[0])

        # EXEC from start
        exec_token = SMToken(
            target=0, addr=256, op=MemOp.EXEC,
            flags=None, data=None, ret=None
        )

        def do_exec():
            yield sm.input_store.put(exec_token)

        env.process(do_exec())
        # Should not crash
        env.run()

        # No tokens should have been injected (truncated)
        assert len(mock_pe_store.items) == 0

    def test_mixed_valid_invalid_packets_stops_at_first_error(self):
        """EXEC stops at first malformed packet, valid packets before it are processed."""
        env = simpy.Environment()
        sm = StructureMemory(env, 0, cell_count=512, tier_boundary=256)

        # Mock system and PE input store
        from unittest.mock import Mock
        mock_pe_store = simpy.Store(env)
        mock_system = Mock()
        sm.system = mock_system

        def mock_send(token):
            yield env.timeout(1)
            yield mock_pe_store.put(token)

        mock_system.send = mock_send

        # Valid token first
        token1 = DyadToken(target=0, offset=5, act_id=1, data=0x1111, port=Port.L)
        flits1 = pack_token(token1)

        # Then invalid
        invalid = [0xFFFF]

        # Pre-load
        sm.t0_store.extend(flits1)
        sm.t0_store.extend(invalid)

        # EXEC from start
        exec_token = SMToken(
            target=0, addr=256, op=MemOp.EXEC,
            flags=None, data=None, ret=None
        )

        def do_exec():
            yield sm.input_store.put(exec_token)

        env.process(do_exec())
        env.run()

        # First valid token should have been injected before error
        assert len(mock_pe_store.items) >= 1
        assert mock_pe_store.items[0].data == 0x1111