"""Tests for code generation (frame-based model).

Tests verify:
- pe-frame-redesign.AC6.1: AssemblyResult includes setup_tokens field
- pe-frame-redesign.AC6.2: Token stream ordering: SM init → IRAM writes → ALLOC → frame slot writes → seeds
- pe-frame-redesign.AC6.3: IRAM write data uses pack_instruction()
- pe-frame-redesign.AC6.6: Seed tokens use act_id, no gen field

Also tests original codegen acceptance criteria:
- Direct mode produces valid PEConfig with correct IRAM contents
- Direct mode produces valid SMConfig with initial cell values
- Direct mode produces seed tokens for const nodes with no incoming edges
- Direct mode PEConfig includes route restrictions matching edge analysis
- Program with no data_defs produces empty SM init section
"""

import pytest

from asm.codegen import generate_direct, generate_tokens, AssemblyResult
from asm.ir import (
    IRGraph,
    IRNode,
    IREdge,
    IRDataDef,
    SystemConfig,
    SourceLoc,
    ResolvedDest,
)
from cm_inst import (
    Instruction, OutputStyle, TokenKind, FrameOp, FrameDest,
    ArithOp, MemOp, Port, RoutingOp
)
from tokens import MonadToken, SMToken, PELocalWriteToken, FrameControlToken
from emu.types import PEConfig, SMConfig
from sm_mod import Presence


class TestFrameDirectMode:
    """Frame-based direct mode code generation."""

    def test_assembly_result_structure(self):
        """AssemblyResult has all required fields including setup_tokens."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph({"&add": node}, system=system)

        result = generate_direct(graph)

        # Verify AssemblyResult structure
        assert isinstance(result, AssemblyResult)
        assert hasattr(result, 'pe_configs')
        assert hasattr(result, 'sm_configs')
        assert hasattr(result, 'seed_tokens')
        assert hasattr(result, 'setup_tokens')
        assert isinstance(result.setup_tokens, list)

    def test_direct_mode_peconfig_structure(self):
        """PEConfig has Instruction IRAM (not ALUInst/SMInst) and frame config fields."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=7,
            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.pe_id == 0
        assert 0 in pe_config.iram

        inst = pe_config.iram[0]
        assert isinstance(inst, Instruction)
        assert inst.opcode == ArithOp.ADD
        assert inst.output == OutputStyle.INHERIT
        assert inst.fref == 7

        # Verify frame config fields
        assert hasattr(pe_config, 'frame_count')
        assert hasattr(pe_config, 'frame_slots')
        assert hasattr(pe_config, 'matchable_offsets')
        assert pe_config.frame_count > 0
        assert pe_config.frame_slots > 0
        assert pe_config.matchable_offsets > 0

    def test_direct_mode_with_data_defs(self):
        """Direct mode produces SMConfig with initial cell values."""
        data_def = IRDataDef(
            name="@val",
            sm_id=0,
            cell_addr=5,
            value=42,
            loc=SourceLoc(1, 1),
        )
        node = IRNode(
            name="&a",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&a": node},
            data_defs=[data_def],
            system=system,
        )

        result = generate_direct(graph)

        # Verify SMConfig
        assert len(result.sm_configs) == 1
        sm_config = result.sm_configs[0]
        assert sm_config.sm_id == 0
        assert 5 in sm_config.initial_cells
        pres, val = sm_config.initial_cells[5]
        assert pres == Presence.FULL
        assert val == 42

    def test_seed_tokens_use_act_id(self):
        """Seed tokens use act_id field, not ctx."""
        seed_node = IRNode(
            name="&seed",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=1,
            act_id=3,
            const=99,
            mode=(OutputStyle.INHERIT, True, 1),
            wide=False,
            fref=0,
            seed=True,
            loc=SourceLoc(2, 1),
        )
        # Monadic destination (not dyadic)
        dest_node = IRNode(
            name="&dest",
            opcode=ArithOp.INC,  # INC is monadic
            pe=0,
            iram_offset=2,
            act_id=3,
            mode=(OutputStyle.INHERIT, False, 1),
            wide=False,
            fref=0,
            loc=SourceLoc(3, 1),
        )
        edge = IREdge(source="&seed", dest="&dest", port=Port.L)
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&seed": seed_node, "&dest": dest_node},
            edges=[edge],
            system=system,
        )

        result = generate_direct(graph)

        # Verify seed tokens
        assert len(result.seed_tokens) == 1
        seed = result.seed_tokens[0]
        assert isinstance(seed, MonadToken)
        assert seed.data == 99
        assert seed.act_id == 3
        assert not hasattr(seed, 'gen')  # No gen field

    def test_const_node_with_no_incoming_edge_is_seed(self):
        """CONST node with no incoming edges produces seed token."""
        const_node = IRNode(
            name="&const",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=2,
            act_id=0,
            const=99,
            mode=(OutputStyle.INHERIT, True, 1),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )
        graph = IRGraph({"&const": 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.act_id == 0
        assert token.data == 99

    def test_route_restrictions(self):
        """Cross-PE edges produce correct allowed_pe_routes."""
        node_pe0 = IRNode(
            name="&a",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            dest_l=ResolvedDest(name="&b", frame_dest=FrameDest(
                target_pe=1, offset=0, act_id=0, port=Port.L, token_kind=TokenKind.DYADIC
            )),
            loc=SourceLoc(1, 1),
        )
        node_pe1 = IRNode(
            name="&b",
            opcode=ArithOp.SUB,
            pe=1,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=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)
        assert 1 in pe1_config.allowed_pe_routes

    def test_no_data_defs_produces_empty_smconfig(self):
        """Program with no data_defs produces SMConfig with no initial cells."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=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) == 1
        sm_config = result.sm_configs[0]
        assert sm_config.sm_id == 0
        assert sm_config.initial_cells is None

    def test_memop_instructions(self):
        """MemOp instructions produce Instruction objects."""
        node = IRNode(
            name="&read",
            opcode=MemOp.READ,
            pe=0,
            iram_offset=0,
            act_id=0,
            sm_id=0,
            mode=(OutputStyle.CHANGE_TAG, False, 1),
            wide=True,
            fref=3,
            loc=SourceLoc(1, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph({"&read": node}, system=system)

        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, Instruction)
        assert inst.opcode == MemOp.READ
        assert inst.output == OutputStyle.CHANGE_TAG
        assert inst.wide == True
        assert inst.fref == 3


class TestTokenStream:
    """Token stream generation and ordering."""

    def test_setup_tokens_include_sm_init(self):
        """Setup tokens include SM init (SMToken) entries."""
        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,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=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)
        sm_tokens = [t for t in result.setup_tokens if isinstance(t, SMToken)]

        assert len(sm_tokens) > 0
        token = sm_tokens[0]
        assert token.target == 0
        assert token.addr == 5
        assert token.op == MemOp.WRITE
        assert token.data == 42

    def test_setup_tokens_include_iram_writes(self):
        """Setup tokens include IRAM writes (PELocalWriteToken with region=0)."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph({"&add": node}, system=system)

        result = generate_direct(graph)
        iram_tokens = [t for t in result.setup_tokens
                       if isinstance(t, PELocalWriteToken) and t.region == 0]

        assert len(iram_tokens) > 0
        token = iram_tokens[0]
        assert token.target == 0
        assert token.region == 0

    def test_setup_tokens_include_alloc(self):
        """Setup tokens include ALLOC (FrameControlToken) entries."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=2,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph({"&add": node}, system=system)

        result = generate_direct(graph)
        alloc_tokens = [t for t in result.setup_tokens if isinstance(t, FrameControlToken)]

        assert len(alloc_tokens) > 0
        token = alloc_tokens[0]
        assert token.op == FrameOp.ALLOC
        assert token.act_id == 2

    def test_generate_tokens_ordering(self):
        """generate_tokens() returns tokens in proper order."""
        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,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )
        seed_node = IRNode(
            name="&const",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=10,
            act_id=0,
            const=99,
            mode=(OutputStyle.INHERIT, True, 1),
            wide=False,
            fref=0,
            seed=True,
            loc=SourceLoc(2, 1),
        )
        edge = IREdge(source="&const", dest="&add", port=Port.L)
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&add": node, "&const": seed_node},
            edges=[edge],
            data_defs=[data_def],
            system=system,
        )

        tokens = generate_tokens(graph)

        # Verify that setup tokens come before seed tokens
        sm_indices = [i for i, t in enumerate(tokens) if isinstance(t, SMToken)]
        iram_indices = [i for i, t in enumerate(tokens)
                       if isinstance(t, PELocalWriteToken) and t.region == 0]
        seed_indices = [i for i, t in enumerate(tokens) if isinstance(t, (MonadToken))]

        if sm_indices and iram_indices:
            assert max(sm_indices) < min(iram_indices)
        if iram_indices and seed_indices:
            assert max(iram_indices) < min(seed_indices)


class TestEdgeCases:
    """Edge case tests."""

    def test_multiple_data_defs_same_sm(self):
        """Multiple data_defs targeting same SM are merged."""
        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 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."""
        source_node = IRNode(
            name="&src",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )
        const_node = IRNode(
            name="&const",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=1,
            act_id=0,
            const=5,
            mode=(OutputStyle.INHERIT, True, 1),
            wide=False,
            fref=0,
            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

    def test_multiactivation_alloc_tokens(self):
        """Multiple activations generate multiple ALLOC tokens."""
        node1 = IRNode(
            name="&a",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=1,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )
        node2 = IRNode(
            name="&b",
            opcode=ArithOp.SUB,
            pe=0,
            iram_offset=1,
            act_id=2,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(2, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&a": node1, "&b": node2},
            system=system,
        )

        result = generate_direct(graph)
        alloc_tokens = [t for t in result.setup_tokens if isinstance(t, FrameControlToken)]

        # Should have 2 ALLOC tokens (one for act_id=1, one for act_id=2)
        assert len(alloc_tokens) == 2
        assert all(t.op == FrameOp.ALLOC for t in alloc_tokens)
        assert {t.act_id for t in alloc_tokens} == {1, 2}