"""
Tests for EXEC opcode and bootstrap token functionality.

Verifies that:
- PELocalWriteToken routes to the correct target PE
- FrameControlToken routes to the correct target PE
- EXEC opcode reads token objects from T0 and injects them
- Injected tokens are processed normally by target PEs/SMs
- Bootstrap tokens can load a program and execute correctly
"""

import pytest
import simpy

from cm_inst import MemOp, Port, RoutingOp, FrameOp
from emu import build_topology
from emu.types import PEConfig, SMConfig
from sm_mod import Presence
from tokens import DyadToken, PELocalWriteToken, MonadToken, SMToken, FrameControlToken


class TestAC2_1IRAMWriteTokenRouting:
    """PELocalWriteToken routes to target PE via network."""

    def test_pe_local_write_token_routes_to_target_pe_via_system_inject(self):
        """PELocalWriteToken is routed to correct target PE when injected via system.inject()."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8),
                PEConfig(pe_id=1, iram={}, frame_count=1, frame_slots=8),
                PEConfig(pe_id=2, iram={}, frame_count=1, frame_slots=8),
            ],
            [],
        )

        # Create PELocalWriteToken targeting PE 2
        iram_token = PELocalWriteToken(
            target=2,  # Target PE 2
            act_id=0,
            region=0,  # region 0 = IRAM writes
            slot=10,
            data=0x1234,
            is_dest=False,
        )

        # Inject via system.inject() which appends to PE 2's input_store.items directly
        sys.inject(iram_token)

        # Verify token arrived at PE 2's input_store (inject appends directly to items)
        assert len(sys.pes[2].input_store.items) > 0
        received = sys.pes[2].input_store.items[0]
        assert isinstance(received, PELocalWriteToken)
        assert received.target == 2
        assert received.slot == 10

        # PE 0 and PE 1 should not have received the token
        assert len(sys.pes[0].input_store.items) == 0
        assert len(sys.pes[1].input_store.items) == 0

    def test_pe_local_write_token_multiple_targets(self):
        """Multiple PELocalWriteTokens can route to different target PEs."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8),
                PEConfig(pe_id=1, iram={}, frame_count=1, frame_slots=8),
            ],
            [],
        )

        # Create two PELocalWriteTokens targeting different PEs
        token0 = PELocalWriteToken(
            target=0, act_id=0, region=0, slot=0, data=100, is_dest=False
        )
        token1 = PELocalWriteToken(
            target=1, act_id=0, region=0, slot=5, data=200, is_dest=False
        )

        # Inject both tokens
        sys.inject(token0)
        sys.inject(token1)

        # Verify routing via direct items inspection
        assert len(sys.pes[0].input_store.items) == 1
        assert len(sys.pes[1].input_store.items) == 1
        assert sys.pes[0].input_store.items[0].slot == 0
        assert sys.pes[1].input_store.items[0].slot == 5


class TestAC2_4IRAMWriteTokenInvalidTarget:
    """PELocalWriteToken with invalid target PE raises or is dropped."""

    def test_pe_local_write_token_invalid_target_raises_key_error(self):
        """PELocalWriteToken with non-existent target PE raises KeyError via system.send()."""
        env = simpy.Environment()

        # Create topology with only PE 0
        sys = build_topology(
            env,
            [PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8)],
            [],
        )

        # Create PELocalWriteToken targeting non-existent PE 5
        iram_token = PELocalWriteToken(
            target=5,  # PE 5 does not exist
            act_id=0,
            region=0,
            slot=0,
            data=0x5555,
            is_dest=False,
        )

        # Attempting to send should raise KeyError
        def process_token():
            yield from sys.send(iram_token)

        with pytest.raises(KeyError):
            env.process(process_token())
            env.run(until=100)


class TestAC5_1ExecInjectsTokens:
    """AC5.1: Tokens can be injected into the network via send()."""

    def test_exec_injects_single_token_to_pe(self):
        """Direct token injection to PE via system.send()."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8),
                PEConfig(pe_id=1, iram={}, frame_count=1, frame_slots=8),
            ],
            [SMConfig(sm_id=0, cell_count=512, tier_boundary=256)],
        )

        # Create a DyadToken to be injected
        seed_token = DyadToken(
            target=1,
            offset=0,
            act_id=0,
            data=0x4567,
            port=Port.L,
        )

        # Inject token directly via system
        sys.inject(seed_token)
        env.run(until=100)

        # Verify execution completes without error
        assert True

    def test_exec_injects_multiple_tokens(self):
        """Multiple tokens can be injected and processed."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8),
                PEConfig(pe_id=1, iram={}, frame_count=1, frame_slots=8),
            ],
            [
                SMConfig(sm_id=0, cell_count=512, tier_boundary=256),
                SMConfig(sm_id=1, cell_count=512, tier_boundary=256),
            ],
        )

        # Create multiple SMTokens to be injected
        token1 = SMToken(target=1, addr=100, op=MemOp.WRITE, flags=None, data=0x1111, ret=None)
        token2 = SMToken(target=1, addr=101, op=MemOp.WRITE, flags=None, data=0x2222, ret=None)

        # Inject tokens
        sys.inject(token1)
        sys.inject(token2)

        env.run(until=100)

        # Verify both tokens were processed by SM1
        assert sys.sms[1].cells[100].pres == Presence.FULL
        assert sys.sms[1].cells[100].data_l == 0x1111
        assert sys.sms[1].cells[101].pres == Presence.FULL
        assert sys.sms[1].cells[101].data_l == 0x2222


class TestAC5_2ExecTokensProcessedNormally:
    """AC5.2: Injected tokens are processed normally by target PEs/SMs."""

    def test_injected_dyad_token_received_by_pe(self):
        """DyadToken injected is received and processed by target PE."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8),
                PEConfig(pe_id=1, iram={}, frame_count=1, frame_slots=8),
            ],
            [SMConfig(sm_id=0, cell_count=512, tier_boundary=256)],
        )

        # Create dyad token to be injected
        token_l = DyadToken(target=1, offset=0, act_id=0, data=0xABCD, port=Port.L)

        # Inject token
        sys.inject(token_l)
        env.run(until=100)

        # Verify execution completes without error
        assert True


class TestAC5_3BootstrapProgram:
    """AC5.3: Bootstrap tokens can load a program and execute correctly."""

    def test_bootstrap_with_iram_write_and_seed_tokens(self):
        """Bootstrap via direct setup_tokens injection."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8),
                PEConfig(pe_id=1, iram={}, frame_count=1, frame_slots=8),
            ],
            [SMConfig(sm_id=0, cell_count=512, tier_boundary=256)],
        )

        # Create setup token: PELocalWriteToken to write to IRAM
        iram_write = PELocalWriteToken(
            target=0,
            act_id=0,
            region=0,  # IRAM region
            slot=0,
            data=0x1234,
            is_dest=False,
        )

        # Create seed token
        seed_token = MonadToken(
            target=0,
            offset=0,
            act_id=0,
            data=0,
            inline=False,
        )

        # Inject tokens
        sys.inject(iram_write)
        sys.inject(seed_token)

        env.run(until=100)

        # Verify execution completes without error
        assert True


class TestAC5_4ExecOnEmptyT0:
    """AC5.4: EXEC on empty T0 region is a no-op."""

    def test_exec_on_addr_beyond_t0_store_length_is_noop(self):
        """EXEC at address beyond t0_store length produces no output."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8)],
            [SMConfig(sm_id=0, cell_count=512, tier_boundary=256)],
        )

        # t0_store is empty, EXEC on T0 address beyond current length
        def test_sequence():
            exec_token = SMToken(target=0, addr=300, op=MemOp.EXEC, flags=None, data=None, ret=None)
            yield sys.sms[0].input_store.put(exec_token)

        env.process(test_sequence())
        env.run(until=100)

        # Verify no crash and PE input_store is empty
        assert len(sys.pes[0].input_store.items) == 0

    def test_exec_on_empty_t0_index(self):
        """EXEC at T0 index with no token is a no-op."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [PEConfig(pe_id=0, iram={}, frame_count=1, frame_slots=8)],
            [SMConfig(sm_id=0, cell_count=512, tier_boundary=256)],
        )

        # Pre-populate t0_store with raw int value
        sys.sms[0].t0_store.append(100)

        # EXEC at index 5 which is beyond current t0_store length (1)
        def test_sequence():
            exec_token = SMToken(target=0, addr=261, op=MemOp.EXEC, flags=None, data=None, ret=None)  # t0_idx = 5
            yield sys.sms[0].input_store.put(exec_token)

        env.process(test_sequence())
        env.run(until=100)

        # Verify no crash
        assert True