"""
Tests for initialization API smoke tests.

Verifies:
- or1-emu.AC5.1: IRAM initialization — PE has expected instructions at expected offsets
- or1-emu.AC5.2: SM cell initialization — SM cells match config
- or1-emu.AC5.3: Token injection — inject() routes tokens to correct stores by type
"""

import simpy

from cm_inst import ALUInst, Addr, ArithOp, CfgOp, MemOp, Port, RoutingOp, SMInst
from emu import build_topology, PEConfig, SMConfig
from sm_mod import Presence
from tokens import CMToken, CfgToken, DyadToken, LoadInstToken, MonadToken, SMToken


class TestAC51IRAMInitialization:
    """Test AC5.1: IRAM initialization"""

    def test_iram_contains_instructions_at_configured_offsets(self):
        """IRAM contains ALUInst at offsets specified in PEConfig."""
        env = simpy.Environment()

        pe_iram = {
            0: ALUInst(
                op=ArithOp.ADD,
                dest_l=None,
                dest_r=None,
                const=None,
            ),
            5: ALUInst(
                op=RoutingOp.CONST,
                dest_l=None,
                dest_r=None,
                const=99,
            ),
        }

        sys = build_topology(
            env,
            [PEConfig(pe_id=0, iram=pe_iram)],
            [],
        )

        # Verify instructions are at expected offsets
        assert 0 in sys.pes[0].iram
        assert sys.pes[0].iram[0].op == ArithOp.ADD
        assert 5 in sys.pes[0].iram
        assert sys.pes[0].iram[5].op == RoutingOp.CONST
        assert sys.pes[0].iram[5].const == 99

    def test_iram_does_not_contain_uninitialized_offsets(self):
        """IRAM does not contain offsets not specified in config."""
        env = simpy.Environment()

        pe_iram = {
            0: ALUInst(op=ArithOp.ADD, dest_l=None, dest_r=None, const=None),
            5: ALUInst(op=RoutingOp.CONST, dest_l=None, dest_r=None, const=99),
        }

        sys = build_topology(
            env,
            [PEConfig(pe_id=0, iram=pe_iram)],
            [],
        )

        # Verify uninitialized offsets are NOT in IRAM
        assert 3 not in sys.pes[0].iram
        assert 10 not in sys.pes[0].iram


class TestAC52SMCellInitialization:
    """Test AC5.2: SM cell initialization"""

    def test_sm_cells_initialized_with_presence_and_data(self):
        """SM cells initialized via config contain expected presence and data."""
        env = simpy.Environment()

        sm_config = SMConfig(
            sm_id=0,
            cell_count=512,
            initial_cells={
                0: (Presence.FULL, 42),
                10: (Presence.RESERVED, None),
            },
        )

        sys = build_topology(
            env,
            [],
            [sm_config],
        )

        # Verify cell 0: FULL with data 42
        assert sys.sms[0].cells[0].pres == Presence.FULL
        assert sys.sms[0].cells[0].data_l == 42

        # Verify cell 10: RESERVED with no data
        assert sys.sms[0].cells[10].pres == Presence.RESERVED
        assert sys.sms[0].cells[10].data_l is None

    def test_uninitialized_sm_cells_are_empty(self):
        """SM cells not in initial_cells config are EMPTY."""
        env = simpy.Environment()

        sm_config = SMConfig(
            sm_id=0,
            cell_count=512,
            initial_cells={
                0: (Presence.FULL, 42),
            },
        )

        sys = build_topology(
            env,
            [],
            [sm_config],
        )

        # Verify uninitialized cells are EMPTY
        assert sys.sms[0].cells[1].pres == Presence.EMPTY
        assert sys.sms[0].cells[1].data_l is None
        assert sys.sms[0].cells[100].pres == Presence.EMPTY
        assert sys.sms[0].cells[100].data_l is None


class TestAC53TokenInjection:
    """Test AC5.3: Token injection API"""

    def test_inject_monad_token_to_pe(self):
        """inject() delivers MonadToken to correct PE's input_store."""
        env = simpy.Environment()

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

        # Create and inject token to PE0
        token = MonadToken(
            target=0,
            offset=0,
            ctx=0,
            data=42,
            inline=False,
        )

        sys.inject(token)

        # Verify token is in PE0's input_store
        assert len(sys.pes[0].input_store.items) == 1
        assert sys.pes[0].input_store.items[0] == token

    def test_inject_multiple_tokens_to_correct_pes(self):
        """inject() places tokens in correct PE based on target."""
        env = simpy.Environment()

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

        # Inject token to PE0
        token0 = MonadToken(target=0, offset=0, ctx=0, data=10, inline=False)
        sys.inject(token0)

        # Inject token to PE1
        token1 = MonadToken(target=1, offset=0, ctx=0, data=20, inline=False)
        sys.inject(token1)

        # Verify tokens arrived at correct PEs
        assert len(sys.pes[0].input_store.items) == 1
        assert sys.pes[0].input_store.items[0].data == 10

        assert len(sys.pes[1].input_store.items) == 1
        assert sys.pes[1].input_store.items[0].data == 20

    def test_inject_sm_token_to_sm(self):
        """inject() delivers SMToken to correct SM's input_store."""
        env = simpy.Environment()

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

        # Create and inject token to SM0
        token = SMToken(
            target=0,
            addr=5,
            op=MemOp.READ,
            flags=None,
            data=None,
            ret=CMToken(target=0, offset=0, ctx=0, data=0),
        )

        sys.inject(token)

        # Verify token is in SM0's input_store
        assert len(sys.sms[0].input_store.items) == 1
        assert sys.sms[0].input_store.items[0] == token

    def test_inject_sm_multiple_tokens_to_correct_sms(self):
        """inject() routes SMTokens to correct SM based on token.target."""
        env = simpy.Environment()

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

        # Inject token to SM0
        token0 = SMToken(
            target=0,
            addr=10,
            op=MemOp.WRITE,
            flags=None,
            data=42,
            ret=None,
        )
        sys.inject(token0)

        # Inject token to SM1
        token1 = SMToken(
            target=1,
            addr=20,
            op=MemOp.READ,
            flags=None,
            data=None,
            ret=CMToken(target=0, offset=0, ctx=0, data=0),
        )
        sys.inject(token1)

        # Verify tokens arrived at correct SMs
        assert len(sys.sms[0].input_store.items) == 1
        assert sys.sms[0].input_store.items[0].addr == 10

        assert len(sys.sms[1].input_store.items) == 1
        assert sys.sms[1].input_store.items[0].addr == 20


class TestAC51GenCounterInitialization:
    """Test AC5.1 extended: gen_counter initialization"""

    def test_gen_counters_initialized_from_config(self):
        """PEConfig with gen_counters list initializes PE's gen_counters."""
        env = simpy.Environment()

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

        # Verify gen_counters match config
        assert sys.pes[0].gen_counters == [1, 0, 2, 3]

    def test_gen_counters_default_to_zero_when_none(self):
        """PEConfig with gen_counters=None (default) initializes all to 0."""
        env = simpy.Environment()

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

        # Verify all gen_counters are 0 (ctx_slots default is 16)
        assert sys.pes[0].gen_counters == [0] * 16

    def test_gen_counters_with_custom_ctx_slots(self):
        """gen_counters list length matches ctx_slots."""
        env = simpy.Environment()

        sys = build_topology(
            env,
            [PEConfig(pe_id=0, iram={}, ctx_slots=8, gen_counters=[1, 2, 3, 4, 5, 6, 7, 8])],
            [],
        )

        # Verify gen_counters match provided list
        assert sys.pes[0].gen_counters == [1, 2, 3, 4, 5, 6, 7, 8]
        assert len(sys.pes[0].gen_counters) == 8


class TestAC61E2EConstFedsAdd:
    """Test AC6.1: CONST on PE0 feeds ADD on PE1"""

    def test_const_feeds_add(self):
        """CONST on PE0 emits tokens that arrive at PE1, trigger ADD, produce correct result."""
        env = simpy.Environment()

        # PE0 IRAM: offset 0 = CONST(7), offset 1 = CONST(3)
        pe0_iram = {
            0: ALUInst(
                op=RoutingOp.CONST,
                dest_l=Addr(a=0, port=Port.L, pe=1),
                dest_r=None,
                const=7,
            ),
            1: ALUInst(
                op=RoutingOp.CONST,
                dest_l=Addr(a=0, port=Port.R, pe=1),
                dest_r=None,
                const=3,
            ),
        }

        # PE1 IRAM: offset 0 = ADD routing to PE2 (collector)
        pe1_iram = {
            0: ALUInst(
                op=ArithOp.ADD,
                dest_l=Addr(a=0, port=Port.L, pe=2),
                dest_r=None,
                const=None,
            ),
        }

        # Build topology with 3 PEs (PE2 has no IRAM, acts as collector)
        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram=pe0_iram),
                PEConfig(pe_id=1, iram=pe1_iram),
                PEConfig(pe_id=2, iram={}),
            ],
            [],
        )

        # Set up PE1's routing to direct output to a collector store (not PE2's input_store)
        # This simulates a sink where results are collected without being consumed
        collector_store = simpy.Store(env, capacity=100)
        sys.pes[1].route_table[2] = collector_store

        # Inject tokens via SimPy process
        def injector():
            yield sys.pes[0].input_store.put(MonadToken(target=0, offset=0, ctx=0, data=0, inline=False))
            yield sys.pes[0].input_store.put(MonadToken(target=0, offset=1, ctx=0, data=0, inline=False))

        env.process(injector())

        # Run simulation until quiescence
        env.run(until=1000)

        # Verify collector receives exactly one token with data=10 (7+3)
        assert len(collector_store.items) == 1
        result_token = collector_store.items[0]
        assert result_token.data == 10


class TestAC62E2ESMRoundTrip:
    """Test AC6.2: SM round-trip (PE writes, PE reads)"""

    def test_sm_round_trip(self):
        """PE writes to SM, another PE reads from SM, receives correct data."""
        env = simpy.Environment()

        # PE0 IRAM: offset 0 = SM WRITE, offset 1 = SM READ
        pe0_iram = {
            0: SMInst(op=MemOp.WRITE, sm_id=0, const=0),  # Write to cell 0
            1: SMInst(
                op=MemOp.READ,
                sm_id=0,
                const=0,  # Read from cell 0
                ret=Addr(a=0, port=Port.L, pe=1),
            ),
        }

        # PE1 IRAM: offset 0 = PASS (to pass through SM READ result)
        pe1_iram = {
            0: ALUInst(
                op=RoutingOp.PASS,
                dest_l=Addr(a=0, port=Port.L, pe=1),
                dest_r=None,
                const=None,
            ),
        }

        # Build topology
        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram=pe0_iram),
                PEConfig(pe_id=1, iram=pe1_iram),
            ],
            [
                # SM0: cell 0 starts EMPTY
                SMConfig(sm_id=0, cell_count=512, initial_cells={}),
            ],
        )

        # Set up collector for PE1's output
        collector_store = simpy.Store(env, capacity=100)
        sys.pes[1].route_table[1] = collector_store

        # Inject tokens via SimPy process (FIFO order ensures WRITE before READ)
        def injector():
            yield sys.pes[0].input_store.put(MonadToken(target=0, offset=0, ctx=0, data=42, inline=False))
            yield sys.pes[0].input_store.put(MonadToken(target=0, offset=1, ctx=0, data=0, inline=False))

        env.process(injector())

        # Run simulation
        env.run(until=1000)

        # Verify collector receives a token with data=42 (the value written and read back)
        assert len(collector_store.items) >= 1
        result_token = collector_store.items[0]
        assert result_token.data == 42


class TestAC63E2EDualFanout:
    """Test AC6.3: DUAL mode fan-out to two consumers"""

    def test_dual_fanout(self):
        """DUAL mode emits same result to both PE1 and PE2."""
        env = simpy.Environment()

        # PE0 IRAM: offset 0 = PASS with dual destinations
        pe0_iram = {
            0: ALUInst(
                op=RoutingOp.PASS,
                dest_l=Addr(a=0, port=Port.L, pe=1),
                dest_r=Addr(a=0, port=Port.L, pe=2),
                const=None,
            ),
        }

        # Build topology
        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram=pe0_iram),
                PEConfig(pe_id=1, iram={}),  # PE1: collector
                PEConfig(pe_id=2, iram={}),  # PE2: collector
            ],
            [],
        )

        # Set up collectors for PE1 and PE2
        collector_1 = simpy.Store(env, capacity=100)
        collector_2 = simpy.Store(env, capacity=100)
        sys.pes[0].route_table[1] = collector_1
        sys.pes[0].route_table[2] = collector_2

        # Inject token via SimPy process
        def injector():
            yield sys.pes[0].input_store.put(MonadToken(target=0, offset=0, ctx=0, data=99, inline=False))

        env.process(injector())

        # Run simulation
        env.run(until=1000)

        # Verify each collector receives one token with data=99
        assert len(collector_1.items) == 1
        assert collector_1.items[0].data == 99

        assert len(collector_2.items) == 1
        assert collector_2.items[0].data == 99


class TestAC64E2ESwitchRouting:
    """Test AC6.4: SWITCH mode conditional routing"""

    def test_switch_routing_condition_true(self):
        """SWEQ with equal operands routes data to dest_l, trigger to dest_r."""
        env = simpy.Environment()

        # PE0 IRAM: offset 0 = SWEQ
        pe0_iram = {
            0: ALUInst(
                op=RoutingOp.SWEQ,
                dest_l=Addr(a=0, port=Port.L, pe=1),
                dest_r=Addr(a=0, port=Port.L, pe=2),
                const=None,
            ),
        }

        # Build topology with gen_counters initialized for dyadic matching
        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram=pe0_iram, gen_counters=[0, 0, 0, 0]),
                PEConfig(pe_id=1, iram={}),  # PE1: receives data token
                PEConfig(pe_id=2, iram={}),  # PE2: receives inline trigger
            ],
            [],
        )

        # Set up collectors
        collector_1 = simpy.Store(env, capacity=100)
        collector_2 = simpy.Store(env, capacity=100)
        sys.pes[0].route_table[1] = collector_1
        sys.pes[0].route_table[2] = collector_2

        # Inject two DyadTokens with same data (5, 5) via SimPy process
        def injector():
            yield sys.pes[0].input_store.put(
                DyadToken(
                    target=0,
                    offset=0,
                    ctx=0,
                    data=5,
                    port=Port.L,
                    gen=0,
                    wide=False,
                )
            )
            yield sys.pes[0].input_store.put(
                DyadToken(
                    target=0,
                    offset=0,
                    ctx=0,
                    data=5,
                    port=Port.R,
                    gen=0,
                    wide=False,
                )
            )

        env.process(injector())

        # Run simulation
        env.run(until=1000)

        # When bool_out=True (equal): data → dest_l (collector_1), trigger → dest_r (collector_2)
        assert len(collector_1.items) == 1
        data_token = collector_1.items[0]
        assert data_token.data == 5

        assert len(collector_2.items) == 1
        trigger_token = collector_2.items[0]
        assert trigger_token.inline is True
        assert trigger_token.data == 0

    def test_switch_routing_condition_false(self):
        """SWEQ with unequal operands routes data to dest_r, trigger to dest_l."""
        env = simpy.Environment()

        # PE0 IRAM: offset 0 = SWEQ
        pe0_iram = {
            0: ALUInst(
                op=RoutingOp.SWEQ,
                dest_l=Addr(a=0, port=Port.L, pe=1),
                dest_r=Addr(a=0, port=Port.L, pe=2),
                const=None,
            ),
        }

        # Build topology
        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram=pe0_iram, gen_counters=[0, 0, 0, 0]),
                PEConfig(pe_id=1, iram={}),  # PE1: receives inline trigger
                PEConfig(pe_id=2, iram={}),  # PE2: receives data token
            ],
            [],
        )

        # Set up collectors
        collector_1 = simpy.Store(env, capacity=100)
        collector_2 = simpy.Store(env, capacity=100)
        sys.pes[0].route_table[1] = collector_1
        sys.pes[0].route_table[2] = collector_2

        # Inject two DyadTokens with different data (5, 10) via SimPy process
        def injector():
            yield sys.pes[0].input_store.put(
                DyadToken(
                    target=0,
                    offset=0,
                    ctx=0,
                    data=5,
                    port=Port.L,
                    gen=0,
                    wide=False,
                )
            )
            yield sys.pes[0].input_store.put(
                DyadToken(
                    target=0,
                    offset=0,
                    ctx=0,
                    data=10,
                    port=Port.R,
                    gen=0,
                    wide=False,
                )
            )

        env.process(injector())

        # Run simulation
        env.run(until=1000)

        # When bool_out=False (not equal): data → dest_r (collector_2), trigger → dest_l (collector_1)
        assert len(collector_2.items) == 1
        data_token = collector_2.items[0]
        assert data_token.data == 5  # First operand goes to dest_r

        assert len(collector_1.items) == 1
        trigger_token = collector_1.items[0]
        assert trigger_token.inline is True
        assert trigger_token.data == 0


class TestTask5CfgTokenLoadInst:
    """Test CfgToken LOAD_INST handling in PE"""

    def test_cfg_token_load_inst(self):
        """CfgToken with LOAD_INST dynamically loads instructions into PE IRAM."""
        env = simpy.Environment()

        # Build topology: PE0 starts with empty IRAM, PE1 is collector
        sys = build_topology(
            env,
            [
                PEConfig(pe_id=0, iram={}),  # Empty IRAM initially
                PEConfig(pe_id=1, iram={}),  # Collector
            ],
            [],
        )

        # Set up collector for output
        collector_store = simpy.Store(env, capacity=100)
        sys.pes[0].route_table[1] = collector_store

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

        # Create LoadInstToken to load instruction at offset 0
        cfg_token = LoadInstToken(
            target=0,
            addr=0,
            op=CfgOp.LOAD_INST,
            instructions=(const_inst,),
        )

        # Function to send CfgToken then seed token
        def injector():
            # First, send CfgToken to load the instruction
            yield sys.pes[0].input_store.put(cfg_token)
            # Then inject a MonadToken seed to trigger the loaded instruction
            yield sys.pes[0].input_store.put(
                MonadToken(target=0, offset=0, ctx=0, data=0, inline=False)
            )

        env.process(injector())

        # Run simulation
        env.run(until=1000)

        # Verify:
        # 1. The instruction was loaded into IRAM
        assert 0 in sys.pes[0].iram
        assert sys.pes[0].iram[0].op == RoutingOp.CONST
        assert sys.pes[0].iram[0].const == 42

        # 2. The loaded instruction executed and produced output
        assert len(collector_store.items) == 1
        result_token = collector_store.items[0]
        assert result_token.data == 42