"""
Tests for network topology and routing.

Verifies restricted topology functionality for ProcessingElement network.
"""

import simpy

from cm_inst import OutputStyle, ArithOp, Port, RoutingOp, Instruction, FrameDest, TokenKind
from emu.network import build_topology
from emu.types import PEConfig, SMConfig
from tokens import MonadToken, DyadToken


class TestRestrictedTopology:
    """Test restricted topology via PEConfig allowed routes (AC7.6–AC7.7).

    Verifies:
    - AC7.6: build_topology applies route restrictions from PEConfig
    - AC7.7: PEConfig with None routes preserves full-mesh (backward compatibility)
    """

    def test_ac76_restricted_topology_pe_routes(self):
        """AC7.6: build_topology restricts PE routes based on allowed_pe_routes."""
        env = simpy.Environment()

        # Create 3 PEs but restrict PE 0 to only route to PE 1
        pe0_config = PEConfig(pe_id=0, iram={}, allowed_pe_routes={1})
        pe1_config = PEConfig(pe_id=1, iram={})
        pe2_config = PEConfig(pe_id=2, iram={})

        sys = build_topology(env, [pe0_config, pe1_config, pe2_config], [])

        # PE 0 should only have PE 1 in its route_table
        pe0 = sys.pes[0]
        assert set(pe0.route_table.keys()) == {1}

    def test_ac76_restricted_topology_sm_routes(self):
        """AC7.6: build_topology restricts SM routes based on allowed_sm_routes."""
        env = simpy.Environment()

        # Create PE 0 restricted to SM 0 only (not SM 1)
        pe0_config = PEConfig(pe_id=0, iram={}, allowed_sm_routes={0})
        sm0_config = SMConfig(sm_id=0)
        sm1_config = SMConfig(sm_id=1)

        sys = build_topology(env, [pe0_config], [sm0_config, sm1_config])

        # PE 0 should only have SM 0 in its sm_routes
        pe0 = sys.pes[0]
        assert set(pe0.sm_routes.keys()) == {0}

    def test_ac76_restricted_topology_both_pe_and_sm(self):
        """AC7.6: build_topology applies both PE and SM route restrictions."""
        env = simpy.Environment()

        # Create PE 0 restricted to PE 1 and SM 0 only
        pe0_config = PEConfig(
            pe_id=0,
            iram={},
            allowed_pe_routes={1},
            allowed_sm_routes={0}
        )
        pe1_config = PEConfig(pe_id=1, iram={})
        pe2_config = PEConfig(pe_id=2, iram={})
        sm0_config = SMConfig(sm_id=0)
        sm1_config = SMConfig(sm_id=1)

        sys = build_topology(
            env,
            [pe0_config, pe1_config, pe2_config],
            [sm0_config, sm1_config]
        )

        # PE 0 should be restricted in both dimensions
        pe0 = sys.pes[0]
        assert set(pe0.route_table.keys()) == {1}
        assert set(pe0.sm_routes.keys()) == {0}

        # PE 1 and PE 2 should have full-mesh (no restrictions)
        pe1 = sys.pes[1]
        assert set(pe1.route_table.keys()) == {0, 1, 2}
        assert set(pe1.sm_routes.keys()) == {0, 1}

        pe2 = sys.pes[2]
        assert set(pe2.route_table.keys()) == {0, 1, 2}
        assert set(pe2.sm_routes.keys()) == {0, 1}

    def test_ac77_none_routes_preserves_full_mesh(self):
        """AC7.7: PEConfig with None routes preserves full-mesh topology (backward compat)."""
        env = simpy.Environment()

        # Create 3 PEs with no route restrictions (None)
        pe0_config = PEConfig(pe_id=0, iram={})  # allowed_pe_routes=None, allowed_sm_routes=None
        pe1_config = PEConfig(pe_id=1, iram={})
        pe2_config = PEConfig(pe_id=2, iram={})
        sm0_config = SMConfig(sm_id=0)
        sm1_config = SMConfig(sm_id=1)

        sys = build_topology(
            env,
            [pe0_config, pe1_config, pe2_config],
            [sm0_config, sm1_config]
        )

        # All PEs should have full-mesh routes
        for pe_id in [0, 1, 2]:
            pe = sys.pes[pe_id]
            assert set(pe.route_table.keys()) == {0, 1, 2}
            assert set(pe.sm_routes.keys()) == {0, 1}