"""Tests for code generation.

Tests verify:
- or1-asm.AC8.1: Direct mode produces valid PEConfig with correct IRAM contents
- or1-asm.AC8.2: Direct mode produces valid SMConfig with initial cell values
- or1-asm.AC8.3: Direct mode produces seed MonadTokens for const nodes with no incoming edges
- or1-asm.AC8.4: Direct mode PEConfig includes route restrictions matching edge analysis
- or1-asm.AC8.5: Token stream mode emits SM init tokens before ROUTE_SET tokens
- or1-asm.AC8.6: Token stream mode emits ROUTE_SET tokens before LOAD_INST tokens
- or1-asm.AC8.7: Token stream mode emits LOAD_INST tokens before seed tokens
- or1-asm.AC8.8: Token stream mode produces valid tokens consumable by emulator
- or1-asm.AC8.9: Program with no data_defs produces empty SM init section
- or1-asm.AC8.10: Single PE program produces ROUTE_SET with only self-routes
"""

from asm.codegen import generate_direct, generate_tokens, AssemblyResult
from asm.ir import (
    IRGraph,
    IRNode,
    IREdge,
    IRDataDef,
    SystemConfig,
    SourceLoc,
    ResolvedDest,
)
from cm_inst import ALUInst, Addr, ArithOp, CfgOp, MemOp, Port, RoutingOp, SMInst
from tokens import CfgToken, LoadInstToken, MonadToken, RouteSetToken, SMToken
from emu.types import PEConfig, SMConfig
from sm_mod import Presence


class TestDirectMode:
    """AC8.1, AC8.2, AC8.3, AC8.4: Direct mode code generation."""

    def test_ac81_simple_alu_instructions(self):
        """AC8.1: Two ALU nodes on PE0 produce PEConfig with correct IRAM.

        Tests that:
        - ALU instructions are correctly converted to ALUInst
        - They are placed in IRAM at assigned offsets
        """
        # Create two simple dyadic ALU nodes
        add_node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            ctx=0,
            loc=SourceLoc(1, 1),
        )
        sub_node = IRNode(
            name="&sub",
            opcode=ArithOp.SUB,
            pe=0,
            iram_offset=1,
            ctx=0,
            loc=SourceLoc(2, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&add": add_node, "&sub": sub_node},
            system=system,
        )

        result = generate_direct(graph)

        assert len(result.pe_configs) == 1
        pe_config = result.pe_configs[0]
        assert pe_config.pe_id == 0
        assert len(pe_config.iram) == 2
        assert 0 in pe_config.iram
        assert 1 in pe_config.iram

        # Check the instruction types
        inst_0 = pe_config.iram[0]
        inst_1 = pe_config.iram[1]
        assert isinstance(inst_0, ALUInst)  # Is an ALUInst
        assert isinstance(inst_1, ALUInst)  # Is an ALUInst
        assert inst_0.op == ArithOp.ADD
        assert inst_1.op == ArithOp.SUB

    def test_ac82_data_defs_to_smconfig(self):
        """AC8.2: Data definitions produce SMConfig with initial cell values.

        Tests that:
        - Data defs with SM placement are converted to SMConfig
        - initial_cells dict contains correct (Presence.FULL, value) tuples
        """
        data_def = IRDataDef(
            name="@val",
            sm_id=0,
            cell_addr=5,
            value=42,
            loc=SourceLoc(1, 1),
        )
        graph = IRGraph({}, data_defs=[data_def], system=SystemConfig(1, 1))

        result = generate_direct(graph)

        assert len(result.sm_configs) == 1
        sm_config = result.sm_configs[0]
        assert sm_config.sm_id == 0
        assert sm_config.initial_cells is not None
        assert 5 in sm_config.initial_cells
        pres, val = sm_config.initial_cells[5]
        assert pres == Presence.FULL
        assert val == 42

    def test_ac83_const_node_seed_token(self):
        """AC8.3: CONST node with no incoming edges produces seed MonadToken.

        Tests that:
        - CONST nodes are detected
        - Nodes with no incoming edges are marked as seeds
        - MonadToken has correct target PE, offset, ctx, data
        """
        const_node = IRNode(
            name="&seed",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=2,
            ctx=0,
            const=99,
            loc=SourceLoc(1, 1),
        )
        graph = IRGraph({"&seed": const_node}, system=SystemConfig(1, 1))

        result = generate_direct(graph)

        assert len(result.seed_tokens) == 1
        token = result.seed_tokens[0]
        assert isinstance(token, MonadToken)
        assert token.target == 0
        assert token.offset == 2
        assert token.ctx == 0
        assert token.data == 99
        assert token.inline == False

    def test_ac84_route_restrictions(self):
        """AC8.4: Cross-PE edges produce correct allowed_pe_routes.

        Tests that:
        - Edges from PE0 to PE1 add PE1 to PE0's allowed_pe_routes
        - Self-routes are always included
        """
        # PE0 node connecting to PE1 node
        node_pe0 = IRNode(
            name="&a",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            ctx=0,
            dest_l=ResolvedDest(
                name="&b",
                addr=Addr(a=0, port=Port.L, pe=1),
            ),
            loc=SourceLoc(1, 1),
        )
        node_pe1 = IRNode(
            name="&b",
            opcode=ArithOp.ADD,
            pe=1,
            iram_offset=0,
            ctx=0,
            loc=SourceLoc(2, 1),
        )
        edge = IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(1, 1))
        system = SystemConfig(pe_count=2, sm_count=1)
        graph = IRGraph(
            {"&a": node_pe0, "&b": node_pe1},
            edges=[edge],
            system=system,
        )

        result = generate_direct(graph)

        assert len(result.pe_configs) == 2
        pe0_config = next(c for c in result.pe_configs if c.pe_id == 0)
        pe1_config = next(c for c in result.pe_configs if c.pe_id == 1)

        # PE0 should have routes to {0, 1}
        assert 0 in pe0_config.allowed_pe_routes
        assert 1 in pe0_config.allowed_pe_routes

        # PE1 should have route to {1} (self only, no incoming cross-PE edges)
        assert 1 in pe1_config.allowed_pe_routes

    def test_sm_instructions_in_iram(self):
        """Verify SMInst objects are correctly created and placed in IRAM.

        Tests that MemOp instructions produce SMInst in IRAM.
        """
        sm_node = IRNode(
            name="&read",
            opcode=MemOp.READ,
            pe=0,
            iram_offset=0,
            ctx=0,
            sm_id=0,
            const=42,
            dest_l=ResolvedDest(
                name="&out",
                addr=Addr(a=1, port=Port.L, pe=0),
            ),
            loc=SourceLoc(1, 1),
        )
        graph = IRGraph({"&read": sm_node}, system=SystemConfig(1, 1))

        result = generate_direct(graph)

        assert len(result.pe_configs) == 1
        pe_config = result.pe_configs[0]
        assert 0 in pe_config.iram
        inst = pe_config.iram[0]
        assert isinstance(inst, SMInst)  # Is an SMInst
        assert inst.op == MemOp.READ
        assert inst.sm_id == 0
        assert inst.const == 42


class TestTokenStream:
    """AC8.5, AC8.6, AC8.7, AC8.8: Token stream generation and ordering."""

    def test_ac85_ac86_ac87_token_ordering(self):
        """AC8.5-8.7: Tokens are emitted in correct order.

        Tests that:
        - SM init tokens come first
        - ROUTE_SET tokens come next
        - LOAD_INST tokens come next
        - Seed tokens come last
        """
        # Create a multi-PE graph with data_defs
        data_def = IRDataDef(
            name="@val",
            sm_id=0,
            cell_addr=5,
            value=42,
            loc=SourceLoc(1, 1),
        )
        node1 = IRNode(
            name="&a",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            ctx=0,
            loc=SourceLoc(1, 1),
        )
        node2 = IRNode(
            name="&b",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=1,
            ctx=0,
            const=10,
            loc=SourceLoc(2, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&a": node1, "&b": node2},
            data_defs=[data_def],
            system=system,
        )

        tokens = generate_tokens(graph)

        # Find positions of token types
        smtoken_indices = [
            i for i, t in enumerate(tokens)
            if isinstance(t, SMToken)
        ]
        route_set_indices = [
            i for i, t in enumerate(tokens)
            if isinstance(t, CfgToken) and t.op == CfgOp.ROUTE_SET
        ]
        load_inst_indices = [
            i for i, t in enumerate(tokens)
            if isinstance(t, CfgToken) and t.op == CfgOp.LOAD_INST
        ]
        seed_indices = [
            i for i, t in enumerate(tokens) if isinstance(t, MonadToken)
        ]

        # Verify order: SM < ROUTE_SET < LOAD_INST < seed
        assert smtoken_indices, "Should have at least one SM token"
        assert route_set_indices, "Should have at least one ROUTE_SET token"
        assert load_inst_indices, "Should have at least one LOAD_INST token"
        assert seed_indices, "Should have at least one seed token"
        assert max(smtoken_indices) < min(route_set_indices), "SM tokens should come before ROUTE_SET"
        assert max(route_set_indices) < min(load_inst_indices), "ROUTE_SET should come before LOAD_INST"
        assert max(load_inst_indices) < min(seed_indices), "LOAD_INST should come before seed tokens"

    def test_ac88_tokens_are_valid(self):
        """AC8.8: Generated tokens are valid and consumable by emulator.

        Tests that:
        - All tokens have required fields set
        - Token structure matches emulator expectations
        - Tokens can be injected into an emulator System and execution completes
        """
        from emu.network import build_topology
        import simpy

        data_def = IRDataDef(
            name="@val",
            sm_id=0,
            cell_addr=5,
            value=42,
            loc=SourceLoc(1, 1),
        )
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            ctx=0,
            loc=SourceLoc(1, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&add": node},
            data_defs=[data_def],
            system=system,
        )

        result = generate_direct(graph)
        tokens = generate_tokens(graph)

        # Build emulator system from AssemblyResult configs
        env = simpy.Environment()
        emu_system = build_topology(
            env,
            result.pe_configs,
            result.sm_configs,
            fifo_capacity=16,
        )

        # Inject tokens into the system following the sequence:
        # 1. SM init tokens (paired with SM ID)
        # 2. ROUTE_SET and LOAD_INST CfgTokens
        # 3. Seed MonadTokens
        for token in tokens:
            if isinstance(token, SMToken):
                emu_system.inject(token)
            elif isinstance(token, CfgToken):
                emu_system.inject(token)
            elif isinstance(token, MonadToken):
                emu_system.inject(token)

        # Run the simulation for enough steps to complete initialization
        env.run(until=1000)

        # Verify token structure
        for token in tokens:
            if isinstance(token, SMToken):
                assert isinstance(token.target, int)
                assert isinstance(token.addr, int)
                assert isinstance(token.op, MemOp)
                assert token.op == MemOp.WRITE
            elif isinstance(token, RouteSetToken):
                assert isinstance(token.target, int)
                assert isinstance(token.op, CfgOp)
                assert isinstance(token.pe_routes, frozenset)
                assert isinstance(token.sm_routes, frozenset)
            elif isinstance(token, LoadInstToken):
                assert isinstance(token.target, int)
                assert isinstance(token.op, CfgOp)
                assert isinstance(token.instructions, tuple)
            elif isinstance(token, MonadToken):
                assert isinstance(token.target, int)
                assert isinstance(token.offset, int)
                assert isinstance(token.ctx, int)
                assert isinstance(token.data, int)


class TestEdgeCases:
    """AC8.9, AC8.10: Edge cases for code generation."""

    def test_ac89_no_data_defs(self):
        """AC8.9: Program with no data_defs produces no SMConfig or SM tokens.

        Tests that:
        - sm_configs list is empty
        - Token stream has no SMTokens
        """
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            ctx=0,
            loc=SourceLoc(1, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph({"&add": node}, system=system)

        result = generate_direct(graph)
        assert len(result.sm_configs) == 0

        tokens = generate_tokens(graph)
        sm_tokens = [t for t in tokens if isinstance(t, SMToken)]
        assert len(sm_tokens) == 0

    def test_ac810_single_pe_self_route(self):
        """AC8.10: Single-PE program ROUTE_SET contains only self-route.

        Tests that:
        - allowed_pe_routes contains only the PE's own ID
        - Single-PE graph produces ROUTE_SET with pe_routes=[pe_id]
        """
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            ctx=0,
            loc=SourceLoc(1, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph({"&add": node}, system=system)

        result = generate_direct(graph)
        assert len(result.pe_configs) == 1
        pe_config = result.pe_configs[0]
        assert pe_config.allowed_pe_routes == {0}

        tokens = generate_tokens(graph)
        route_set_tokens = [
            t for t in tokens
            if isinstance(t, RouteSetToken)
        ]
        assert len(route_set_tokens) == 1
        token = route_set_tokens[0]
        assert token.pe_routes == frozenset({0})

    def test_multiple_data_defs_same_sm(self):
        """Multiple data_defs targeting same SM produce single SMConfig.

        Tests that:
        - Multiple data_defs for SM0 are merged into single SMConfig
        - initial_cells contains all entries
        """
        data_def1 = IRDataDef(
            name="@val1",
            sm_id=0,
            cell_addr=5,
            value=42,
            loc=SourceLoc(1, 1),
        )
        data_def2 = IRDataDef(
            name="@val2",
            sm_id=0,
            cell_addr=10,
            value=99,
            loc=SourceLoc(2, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph({}, data_defs=[data_def1, data_def2], system=system)

        result = generate_direct(graph)

        assert len(result.sm_configs) == 1
        sm_config = result.sm_configs[0]
        assert sm_config.sm_id == 0
        assert len(sm_config.initial_cells) == 2
        assert sm_config.initial_cells[5] == (Presence.FULL, 42)
        assert sm_config.initial_cells[10] == (Presence.FULL, 99)

    def test_const_node_with_incoming_edge_not_seed(self):
        """CONST node with incoming edge is not a seed token.

        Tests that:
        - Only CONST nodes with NO incoming edges produce seed_tokens
        """
        source_node = IRNode(
            name="&src",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            ctx=0,
            loc=SourceLoc(1, 1),
        )
        const_node = IRNode(
            name="&const",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=1,
            ctx=0,
            const=5,
            loc=SourceLoc(2, 1),
        )
        edge = IREdge(source="&src", dest="&const", port=Port.L, loc=SourceLoc(1, 1))
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&src": source_node, "&const": const_node},
            edges=[edge],
            system=system,
        )

        result = generate_direct(graph)

        # The CONST node has an incoming edge, so it should NOT be a seed
        assert len(result.seed_tokens) == 0


class TestMultiPERouting:
    """Extended tests for multi-PE routing scenarios."""

    def test_multi_pe_route_computation(self):
        """Multi-PE graph with multiple cross-PE edges.

        Tests route computation across multiple PEs with various edge patterns.
        """
        # Create a 3-PE system with cross-PE edges
        node_pe0 = IRNode(
            name="&a",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            ctx=0,
            dest_l=ResolvedDest(
                name="&b",
                addr=Addr(a=0, port=Port.L, pe=1),
            ),
            loc=SourceLoc(1, 1),
        )
        node_pe1 = IRNode(
            name="&b",
            opcode=ArithOp.SUB,
            pe=1,
            iram_offset=0,
            ctx=0,
            dest_l=ResolvedDest(
                name="&c",
                addr=Addr(a=0, port=Port.L, pe=2),
            ),
            loc=SourceLoc(2, 1),
        )
        node_pe2 = IRNode(
            name="&c",
            opcode=ArithOp.INC,
            pe=2,
            iram_offset=0,
            ctx=0,
            loc=SourceLoc(3, 1),
        )
        edge1 = IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(1, 1))
        edge2 = IREdge(source="&b", dest="&c", port=Port.L, loc=SourceLoc(2, 1))
        system = SystemConfig(pe_count=3, sm_count=1)
        graph = IRGraph(
            {"&a": node_pe0, "&b": node_pe1, "&c": node_pe2},
            edges=[edge1, edge2],
            system=system,
        )

        result = generate_direct(graph)

        pe0_config = next(c for c in result.pe_configs if c.pe_id == 0)
        pe1_config = next(c for c in result.pe_configs if c.pe_id == 1)
        pe2_config = next(c for c in result.pe_configs if c.pe_id == 2)

        # PE0 -> PE1
        assert 1 in pe0_config.allowed_pe_routes
        # PE1 -> PE2
        assert 2 in pe1_config.allowed_pe_routes
        # PE2 has no outgoing edges
        assert pe2_config.allowed_pe_routes == {2}