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

Verifies acceptance criteria:
- token-migration.AC2.1: IRAMWriteToken routes to target PE via network (isinstance CMToken)
- token-migration.AC2.4: IRAMWriteToken with invalid target PE raises or is dropped
- token-migration.AC5.1: EXEC reads Token objects from T0 and injects them
- token-migration.AC5.2: Injected tokens are processed normally by target PEs/SMs
- token-migration.AC5.3: EXEC can load a program (IRAM + seed tokens) and execute correctly
- token-migration.AC5.4: EXEC on empty T0 region is a no-op
"""

import pytest
import simpy

from cm_inst import ALUInst, Addr, MemOp, Port, RoutingOp
from emu import build_topology
from emu.types import PEConfig, SMConfig
from sm_mod import Presence
from tokens import DyadToken, IRAMWriteToken, MonadToken, SMToken


class TestAC2_1IRAMWriteTokenRouting:
    """AC2.1: IRAMWriteToken routes to target PE via network (isinstance CMToken)."""

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

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

        # Create IRAMWriteToken targeting PE 2
        inst = ALUInst(
            op=RoutingOp.CONST,
            dest_l=Addr(a=0, port=Port.L, pe=0),
            dest_r=None,
            const=0x1234,
        )
        iram_token = IRAMWriteToken(
            target=2,  # Target PE 2
            offset=10,
            ctx=0,
            data=0,
            instructions=(inst,),
        )

        # 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, IRAMWriteToken)
        assert received.target == 2
        assert received.offset == 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_iram_write_token_multiple_targets(self):
        """Multiple IRAMWriteTokens can route to different target PEs."""
        env = simpy.Environment()

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

        # Create two IRAMWriteTokens targeting different PEs
        inst0 = ALUInst(op=RoutingOp.CONST, dest_l=None, dest_r=None, const=100)
        inst1 = ALUInst(op=RoutingOp.CONST, dest_l=None, dest_r=None, const=200)

        token0 = IRAMWriteToken(target=0, offset=0, ctx=0, data=0, instructions=(inst0,))
        token1 = IRAMWriteToken(target=1, offset=5, ctx=0, data=0, instructions=(inst1,))

        # 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].offset == 0
        assert sys.pes[1].input_store.items[0].offset == 5


class TestAC2_4IRAMWriteTokenInvalidTarget:
    """AC2.4: IRAMWriteToken with invalid target PE raises or is dropped."""

    def test_iram_write_token_invalid_target_raises_key_error(self):
        """IRAMWriteToken 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={})],
            [],
        )

        # Create IRAMWriteToken targeting non-existent PE 5
        inst = ALUInst(op=RoutingOp.CONST, dest_l=None, dest_r=None, const=0x5555)
        iram_token = IRAMWriteToken(
            target=5,  # PE 5 does not exist
            offset=0,
            ctx=0,
            data=0,
            instructions=(inst,),
        )

        # 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: EXEC reads Token objects from T0 and injects them into the network via send()."""

    def test_exec_injects_single_token_to_pe(self):
        """EXEC at T0 address reads a DyadToken and injects it via send() which triggers SimPy events."""
        env = simpy.Environment()

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

        # Create a DyadToken to be injected by EXEC
        seed_token = DyadToken(
            target=1,
            offset=0,
            ctx=0,
            data=0x4567,
            port=Port.L,
            gen=0,
            wide=False,
        )

        # Pre-populate T0 with the token
        sys.sms[0].t0_store.append(seed_token)
        sys.sms[0].system = sys

        def test_sequence():
            # SM0 executes EXEC at T0 address 256 (t0_idx=0)
            exec_token = SMToken(target=0, addr=256, 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 token was injected via send() - it will be consumed by PE1's process
        # and stored in matching_store. The key is that send() triggers the get() event.
        assert sys.pes[1].matching_store[0][0].occupied is True
        assert sys.pes[1].matching_store[0][0].data == 0x4567
        assert sys.pes[1].matching_store[0][0].port == Port.L

    def test_exec_injects_multiple_tokens(self):
        """EXEC at T0 address reads multiple tokens and injects them in order via send().

        Verifies that send() properly wakes up pending get() operations, allowing multiple
        tokens to be delivered in sequence through SimPy's event system.
        """
        env = simpy.Environment()

        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}),
                PEConfig(pe_id=1, iram={}),
            ],
            [
                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 (write operations don't require IRAM)
        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)

        # Pre-populate T0
        sys.sms[0].t0_store.extend([token1, token2])
        sys.sms[0].system = sys

        def test_sequence():
            exec_token = SMToken(target=0, addr=256, 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 both tokens were injected via send() and processed by SM1
        # Both WRITE operations should have updated SM1's cells
        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 by EXEC via send() is received and processed by target PE."""
        env = simpy.Environment()

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

        # Create dyad token to be injected by EXEC
        token_l = DyadToken(target=1, offset=0, ctx=0, data=0xABCD, port=Port.L, gen=0, wide=False)

        # Pre-populate T0
        sys.sms[0].t0_store.append(token_l)
        sys.sms[0].system = sys

        def test_sequence():
            exec_token = SMToken(target=0, addr=256, 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 PE1 received and processed the token via matching_store
        assert sys.pes[1].matching_store[0][0].occupied is True
        assert sys.pes[1].matching_store[0][0].data == 0xABCD
        assert sys.pes[1].matching_store[0][0].port == Port.L


class TestAC5_3BootstrapProgram:
    """AC5.3: EXEC can load a program (IRAM writes + seed tokens) from T0 that executes correctly."""

    def test_bootstrap_with_iram_write_and_seed_tokens(self):
        """Full bootstrap: T0 contains IRAMWriteToken and seed tokens, EXEC loads and runs them.

        Tests the FULL SimPy execution chain:
        - Populate T0 with IRAMWriteToken + seed DyadToken pair
        - Send EXEC SMToken to SM
        - Run env.run()
        - Assert on pe.output_log containing the expected ALU result
        """
        env = simpy.Environment()

        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}),  # PE0 starts empty, will be loaded by bootstrap
                PEConfig(pe_id=1, iram={}),  # PE1 is output receiver
            ],
            [SMConfig(sm_id=0, cell_count=512, tier_boundary=256)],
        )

        # Create instruction to be loaded: CONST(0xABCD) to PE1
        const_inst = ALUInst(
            op=RoutingOp.CONST,
            dest_l=Addr(a=0, port=Port.L, pe=1),
            dest_r=None,
            const=0xABCD,
        )

        # Create IRAMWriteToken to load instruction at offset 0
        iram_write = IRAMWriteToken(
            target=0,
            offset=0,
            ctx=0,
            data=0,
            instructions=(const_inst,),
        )

        # Create seed MonadToken to trigger the loaded instruction at PE0
        seed_token = MonadToken(
            target=0,
            offset=0,
            ctx=0,
            data=0,
            inline=False,
        )

        # Pre-populate T0 with bootstrap sequence
        sys.sms[0].t0_store.append(iram_write)
        sys.sms[0].t0_store.append(seed_token)
        sys.sms[0].system = sys

        def test_sequence():
            # Trigger EXEC to bootstrap
            exec_token = SMToken(target=0, addr=256, 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=200)

        # Verify FULL SimPy execution chain:
        # 1. PE0 should have received and processed IRAMWriteToken
        assert 0 in sys.pes[0].iram, "Instruction not loaded into IRAM by bootstrap"
        assert sys.pes[0].iram[0].op == RoutingOp.CONST
        assert sys.pes[0].iram[0].const == 0xABCD

        # 2. PE0 should have received and processed seed token, producing output
        assert len(sys.pes[0].output_log) > 0, \
            "PE0 did not produce output; seed token may not have triggered IRAM execution"

        # 3. The output should be routed to PE1 with the CONST value
        output_to_pe1 = [t for t in sys.pes[0].output_log if t.target == 1]
        assert len(output_to_pe1) > 0, "PE0 did not route output to PE1"
        assert output_to_pe1[0].data == 0xABCD, f"Expected output data 0xABCD, got {output_to_pe1[0].data}"


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={})],
            [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 tokens were injected (output stores remain unchanged)
        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={})],
            [SMConfig(sm_id=0, cell_count=512, tier_boundary=256)],
        )

        # Pre-populate t0_store with one token
        sys.sms[0].t0_store.append(MonadToken(target=0, offset=0, ctx=0, data=100, inline=False))

        # 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 only the pre-existing token remains in PE0 input_store (no injection happened)
        # The pre-existing token at index 0 should not be re-injected by EXEC at index 5
        assert len(sys.pes[0].input_store.items) == 0