"""
Tests for network routing with frame-based PE.

Verifies pe-frame-redesign.AC1.5:
- Network routing uses isinstance(token, PEToken) for all PE-bound tokens
- CMToken, PELocalWriteToken, FrameControlToken all route to target PE
- SMToken routes to target SM
- build_topology() constructs PE with frame_count, frame_slots, matchable_offsets
- build_topology() loads initial_frames and initial_tag_store
"""

import simpy
import pytest

from emu import build_topology, PEConfig, SMConfig
from emu.types import PEConfig as PECfg
from tokens import (
    CMToken, DyadToken, MonadToken, PELocalWriteToken, FrameControlToken,
    SMToken,
)
from cm_inst import MemOp, Port, FrameOp


class TestPETokenRouting:
    """AC1.5: All PE-bound tokens (CMToken, PELocalWriteToken, FrameControlToken) route to target PE."""

    def test_cmtoken_routes_to_target_pe(self):
        """CMToken (as DyadToken subclass) routes to target PE."""
        env = simpy.Environment()
        pe_configs = [
            PECfg(pe_id=0),
            PECfg(pe_id=1),
        ]
        sm_configs = [SMConfig(sm_id=0)]
        system = build_topology(env, pe_configs, sm_configs)

        # Inject DyadToken to PE 1
        token = DyadToken(
            target=1,
            offset=0,
            act_id=0,
            data=42,
            port=Port.L,
        )
        system.inject(token)

        # Verify token arrived at PE 1's input store
        assert len(system.pes[1].input_store.items) > 0
        assert system.pes[1].input_store.items[0] == token

    def test_pelocal_write_token_routes_to_target_pe(self):
        """PELocalWriteToken routes to target PE."""
        env = simpy.Environment()
        pe_configs = [
            PECfg(pe_id=0),
            PECfg(pe_id=1),
        ]
        sm_configs = [SMConfig(sm_id=0)]
        system = build_topology(env, pe_configs, sm_configs)

        # Inject PELocalWriteToken to PE 1
        token = PELocalWriteToken(
            target=1,
            act_id=0,
            region=0,
            slot=10,
            data=99,
            is_dest=False,
        )
        system.inject(token)

        # Verify token arrived at PE 1's input store
        assert len(system.pes[1].input_store.items) > 0
        assert system.pes[1].input_store.items[0] == token

    def test_frame_control_token_routes_to_target_pe(self):
        """FrameControlToken routes to target PE."""
        env = simpy.Environment()
        pe_configs = [
            PECfg(pe_id=0),
            PECfg(pe_id=1),
        ]
        sm_configs = [SMConfig(sm_id=0)]
        system = build_topology(env, pe_configs, sm_configs)

        # Inject FrameControlToken to PE 1
        token = FrameControlToken(
            target=1,
            act_id=0,
            op=FrameOp.ALLOC,
            payload=0,
        )
        system.inject(token)

        # Verify token arrived at PE 1's input store
        assert len(system.pes[1].input_store.items) > 0
        assert system.pes[1].input_store.items[0] == token

    def test_smtoken_routes_to_target_sm(self):
        """SMToken routes to target SM, not PE."""
        env = simpy.Environment()
        pe_configs = [PECfg(pe_id=0)]
        sm_configs = [
            SMConfig(sm_id=0),
            SMConfig(sm_id=1),
        ]
        system = build_topology(env, pe_configs, sm_configs)

        # Inject SMToken to SM 1
        token = SMToken(
            target=1,
            addr=100,
            op=MemOp.WRITE,
            flags=None,
            data=42,
            ret=None,
        )
        system.inject(token)

        # Verify token arrived at SM 1's input store
        assert len(system.sms[1].input_store.items) > 0
        assert system.sms[1].input_store.items[0] == token
        # And NOT at PE 0
        assert len(system.pes[0].input_store.items) == 0


class TestBuildTopologyFrameConfig:
    """AC1.5: build_topology passes frame config fields to PE constructor."""

    def test_frame_config_params_passed(self):
        """build_topology passes frame_count, frame_slots, matchable_offsets."""
        env = simpy.Environment()
        pe_configs = [
            PECfg(
                pe_id=0,
                frame_count=16,
                frame_slots=128,
                matchable_offsets=16,
            ),
        ]
        sm_configs = [SMConfig(sm_id=0)]
        system = build_topology(env, pe_configs, sm_configs)

        pe = system.pes[0]
        assert pe.frame_count == 16
        assert pe.frame_slots == 128
        assert pe.matchable_offsets == 16

    def test_loads_initial_frames(self):
        """build_topology loads initial_frames into PE."""
        env = simpy.Environment()
        initial_frames = {
            0: [1, 2, 3, 4, 5],
            1: [10, 20, 30, 40, 50],
        }
        pe_configs = [
            PECfg(
                pe_id=0,
                frame_count=8,
                frame_slots=64,
                initial_frames=initial_frames,
            ),
        ]
        sm_configs = [SMConfig(sm_id=0)]
        system = build_topology(env, pe_configs, sm_configs)

        pe = system.pes[0]
        # Frames are initialized with slots filled from the list
        assert pe.frames[0][:5] == [1, 2, 3, 4, 5]
        assert pe.frames[1][:5] == [10, 20, 30, 40, 50]

    def test_loads_initial_tag_store(self):
        """build_topology loads initial_tag_store into PE and removes frames from free_frames."""
        env = simpy.Environment()
        initial_tag_store = {
            0: (2, 0),  # act_id 0 → frame 2, lane 0
            1: (3, 0),  # act_id 1 → frame 3, lane 0
        }
        pe_configs = [
            PECfg(
                pe_id=0,
                frame_count=8,
                initial_tag_store=initial_tag_store,
            ),
        ]
        sm_configs = [SMConfig(sm_id=0)]
        system = build_topology(env, pe_configs, sm_configs)

        pe = system.pes[0]
        assert pe.tag_store[0] == (2, 0)
        assert pe.tag_store[1] == (3, 0)
        # Frames 2 and 3 should be removed from free_frames
        assert 2 not in pe.free_frames
        assert 3 not in pe.free_frames
        # Other frames should still be in free_frames
        assert 0 in pe.free_frames
        assert 1 in pe.free_frames

    def test_initial_frames_and_tag_store_together(self):
        """build_topology correctly initializes both frames and tag_store."""
        env = simpy.Environment()
        initial_frames = {
            0: [1, 2, 3],
            1: [4, 5, 6],
            2: [7, 8, 9],
        }
        initial_tag_store = {
            0: (0, 0),  # act_id 0 uses frame 0, lane 0
            1: (1, 0),  # act_id 1 uses frame 1, lane 0
        }
        pe_configs = [
            PECfg(
                pe_id=0,
                frame_count=4,
                initial_frames=initial_frames,
                initial_tag_store=initial_tag_store,
            ),
        ]
        sm_configs = [SMConfig(sm_id=0)]
        system = build_topology(env, pe_configs, sm_configs)

        pe = system.pes[0]
        # Verify frames loaded (first N slots should match)
        assert pe.frames[0][:3] == [1, 2, 3]
        assert pe.frames[1][:3] == [4, 5, 6]
        assert pe.frames[2][:3] == [7, 8, 9]
        # Verify tag_store loaded
        assert pe.tag_store[0] == (0, 0)
        assert pe.tag_store[1] == (1, 0)
        # Verify free_frames reflects allocations
        assert 0 not in pe.free_frames
        assert 1 not in pe.free_frames
        assert 2 in pe.free_frames
        assert 3 in pe.free_frames


class TestMultiplePEs:
    """Test network routing with multiple PEs and SMs."""

    def test_multiple_pe_network(self):
        """Multiple PEs route tokens correctly."""
        env = simpy.Environment()
        pe_configs = [
            PECfg(pe_id=0),
            PECfg(pe_id=1),
            PECfg(pe_id=2),
        ]
        sm_configs = [SMConfig(sm_id=0)]
        system = build_topology(env, pe_configs, sm_configs)

        # Send tokens to different PEs
        t0 = DyadToken(target=0, offset=0, act_id=0, data=10, port=Port.L)
        t1 = DyadToken(target=1, offset=0, act_id=0, data=20, port=Port.L)
        t2 = DyadToken(target=2, offset=0, act_id=0, data=30, port=Port.L)

        system.inject(t0)
        system.inject(t1)
        system.inject(t2)

        # Verify each PE received correct token
        assert system.pes[0].input_store.items[0].data == 10
        assert system.pes[1].input_store.items[0].data == 20
        assert system.pes[2].input_store.items[0].data == 30

    def test_multiple_sm_network(self):
        """Multiple SMs route tokens correctly."""
        env = simpy.Environment()
        pe_configs = [PECfg(pe_id=0)]
        sm_configs = [
            SMConfig(sm_id=0),
            SMConfig(sm_id=1),
            SMConfig(sm_id=2),
        ]
        system = build_topology(env, pe_configs, sm_configs)

        # Send tokens to different SMs
        t0 = SMToken(target=0, addr=100, op=MemOp.READ, flags=None, data=None, ret=None)
        t1 = SMToken(target=1, addr=100, op=MemOp.READ, flags=None, data=None, ret=None)
        t2 = SMToken(target=2, addr=100, op=MemOp.READ, flags=None, data=None, ret=None)

        system.inject(t0)
        system.inject(t1)
        system.inject(t2)

        # Verify each SM received correct token
        assert system.sms[0].input_store.items[0].target == 0
        assert system.sms[1].input_store.items[0].target == 1
        assert system.sms[2].input_store.items[0].target == 2