"""Tests for frame-based code generation (Phase 6).

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.4: Destination frame slot writes use pack_flit1() with is_dest=True
- pe-frame-redesign.AC6.5: T0 bootstrap data uses pack_token() for packed flits
- pe-frame-redesign.AC6.6: Seed tokens use act_id, no gen field
"""

import pytest

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


class TestTask1BuildIramForPE:
    """Task 1: Rewrite _build_iram_for_pe() for Instruction objects."""

    def test_builds_instruction_objects(self):
        """Produces Instruction objects from IRNodes with mode set."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=5,
            loc=SourceLoc(1, 1),
        )

        iram = _build_iram_for_pe([node], {"&add": node}, [])

        assert len(iram) == 1
        assert 0 in iram
        inst = iram[0]
        assert isinstance(inst, Instruction)
        assert inst.opcode == ArithOp.ADD
        assert inst.output == OutputStyle.INHERIT
        assert inst.has_const == False
        assert inst.dest_count == 2
        assert inst.wide == False
        assert inst.fref == 5

    def test_excludes_seed_nodes(self):
        """Skips seed nodes from IRAM."""
        seed_node = IRNode(
            name="&const",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=10,
            act_id=0,
            mode=(OutputStyle.INHERIT, True, 1),
            wide=False,
            fref=0,
            seed=True,  # Mark as seed
            loc=SourceLoc(1, 1),
        )

        iram = _build_iram_for_pe([seed_node], {"&const": seed_node}, [])

        assert len(iram) == 0

    def test_excludes_unallocated_nodes(self):
        """Skips nodes without iram_offset."""
        node = IRNode(
            name="&unallocated",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=None,  # Not allocated
            act_id=0,
            mode=(OutputStyle.SINK, False, 0),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )

        iram = _build_iram_for_pe([node], {"&unallocated": node}, [])

        assert len(iram) == 0

    def test_excludes_nodes_without_mode(self):
        """Skips nodes without mode (output style) allocation."""
        node = IRNode(
            name="&no_mode",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=5,
            act_id=0,
            mode=None,  # No mode set
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )

        iram = _build_iram_for_pe([node], {"&no_mode": node}, [])

        assert len(iram) == 0

    def test_handles_mem_op_instructions(self):
        """Produces Instruction objects for MemOp opcodes."""
        node = IRNode(
            name="&read",
            opcode=MemOp.READ,
            pe=0,
            iram_offset=1,
            act_id=0,
            sm_id=0,
            mode=(OutputStyle.CHANGE_TAG, False, 1),
            wide=True,
            fref=3,
            loc=SourceLoc(2, 1),
        )

        iram = _build_iram_for_pe([node], {"&read": node}, [])

        assert 1 in iram
        inst = iram[1]
        assert inst.opcode == MemOp.READ
        assert inst.output == OutputStyle.CHANGE_TAG
        assert inst.wide == True
        assert inst.fref == 3


class TestAC63PackInstructionValues:
    """AC6.3: IRAM write data uses pack_instruction() — verify known values."""

    def test_known_instruction_pack_value(self):
        """Verify a known Instruction packs to expected 16-bit value.

        Example from phase design:
        inst = Instruction(opcode=ArithOp.ADD, output=OutputStyle.INHERIT, has_const=False, dest_count=1, wide=False, fref=0)
        Should pack to 0x0000 (all bits zero).
        """
        inst = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.INHERIT,
            has_const=False,
            dest_count=1,
            wide=False,
            fref=0,
        )
        packed = pack_instruction(inst)
        # ADD opcode = 0, INHERIT = 0b000, has_const=False → mode=0b000, wide=0, fref=0
        # [type:1=0][opcode:5=0][mode:3=0][wide:1=0][fref:6=0] = 0x0000
        assert packed == 0x0000

    def test_roundtrip_pack_unpack(self):
        """Verify pack_instruction and unpack_instruction are inverses."""
        orig = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.INHERIT,
            has_const=True,
            dest_count=2,
            wide=True,
            fref=15,
        )
        packed = pack_instruction(orig)
        unpacked = unpack_instruction(packed)
        assert unpacked == orig


class TestTask2FrameSetupTokens:
    """Task 2: Frame setup token generation."""

    def test_ac61_assembly_result_has_setup_tokens(self):
        """AC6.1: AssemblyResult includes setup_tokens field."""
        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 hasattr(result, 'setup_tokens')
        assert isinstance(result.setup_tokens, list)

    def test_ac62_token_ordering(self):
        """AC6.2: Token stream ordering verified.

        Order: SM init → IRAM writes → ALLOC → frame slot writes → seeds
        """
        # Create a simple graph with data def and instruction
        data_def = IRDataDef(
            name="@data",
            sm_id=0,
            cell_addr=10,
            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(2, 1),
        )
        system = SystemConfig(pe_count=1, sm_count=1)
        graph = IRGraph(
            {"&add": node},
            data_defs=[data_def],
            system=system
        )

        result = generate_direct(graph)
        tokens = result.setup_tokens

        # Verify ordering: SM init should come before IRAM writes
        sm_tokens = [t for t in tokens if isinstance(t, SMToken)]
        iram_tokens = [t for t in tokens if isinstance(t, PELocalWriteToken) and t.region == 0]
        alloc_tokens = [t for t in tokens if isinstance(t, FrameControlToken)]

        if sm_tokens and iram_tokens:
            assert tokens.index(sm_tokens[0]) < tokens.index(iram_tokens[0])
        if iram_tokens and alloc_tokens:
            assert tokens.index(iram_tokens[0]) < tokens.index(alloc_tokens[0])

    def test_ac63_iram_write_uses_pack_instruction(self):
        """AC6.3: IRAM write PELocalWriteTokens carry pack_instruction() data."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=5,
            act_id=0,
            mode=(OutputStyle.INHERIT, True, 2),
            wide=False,
            fref=12,
            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]

        # Verify that the data matches packed instruction
        expected_inst = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.INHERIT,
            has_const=True,
            dest_count=2,
            wide=False,
            fref=12,
        )
        expected_data = pack_instruction(expected_inst)
        assert token.data == expected_data

    def test_alloc_tokens_per_activation(self):
        """ALLOC tokens generated for each activation on each PE."""
        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)

        # Verify that PE configs have initial_tag_store with tuple values
        assert len(result.pe_configs) == 1
        pe_cfg = result.pe_configs[0]
        assert pe_cfg.initial_tag_store, "initial_tag_store should not be empty for PE with activations"
        for act_id, val in pe_cfg.initial_tag_store.items():
            assert isinstance(val, tuple) and len(val) == 2, \
                f"initial_tag_store[{act_id}] should be (frame_id, lane) tuple, got {val}"
            frame_id, lane = val
            assert isinstance(frame_id, int), f"frame_id should be int, got {type(frame_id)}"
            assert isinstance(lane, int), f"lane should be int, got {type(lane)}"


class TestTask3SeedTokens:
    """Task 3: Seed token generation with act_id."""

    def test_ac66_seed_tokens_use_act_id(self):
        """AC6.6: Seed tokens use act_id field (not ctx)."""
        seed_node = IRNode(
            name="&const",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=10,
            act_id=5,
            const=42,
            mode=(OutputStyle.INHERIT, True, 1),
            wide=False,
            fref=0,
            seed=True,
            loc=SourceLoc(1, 1),
        )
        edge = IREdge(source="&const", dest="&consumer", port=Port.L)
        consumer = IRNode(
            name="&consumer",
            opcode=ArithOp.ADD,
            pe=1,
            iram_offset=0,
            act_id=3,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(2, 1),
        )
        system = SystemConfig(pe_count=2, sm_count=1)
        graph = IRGraph(
            {"&const": seed_node, "&consumer": consumer},
            edges=[edge],
            system=system
        )

        result = generate_direct(graph)
        seed_tokens = result.seed_tokens

        assert len(seed_tokens) > 0
        token = seed_tokens[0]
        # Verify token uses act_id field
        assert hasattr(token, 'act_id')
        assert token.act_id == 3

    def test_dyadic_seed_token_no_gen(self):
        """DyadToken seed tokens have no gen field."""
        seed_node = IRNode(
            name="&const",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=10,
            act_id=0,
            const=100,
            mode=(OutputStyle.INHERIT, True, 1),
            wide=False,
            fref=0,
            seed=True,
            loc=SourceLoc(1, 1),
        )
        edge = IREdge(source="&const", dest="&dyadic", port=Port.L)
        dyadic_node = IRNode(
            name="&dyadic",
            opcode=ArithOp.ADD,
            pe=1,
            iram_offset=0,
            act_id=1,
            mode=(OutputStyle.INHERIT, False, 2),
            wide=False,
            fref=0,
            loc=SourceLoc(2, 1),
        )
        system = SystemConfig(pe_count=2, sm_count=1)
        graph = IRGraph(
            {"&const": seed_node, "&dyadic": dyadic_node},
            edges=[edge],
            system=system
        )

        result = generate_direct(graph)
        dyadic_tokens = [t for t in result.seed_tokens if isinstance(t, DyadToken)]

        assert len(dyadic_tokens) > 0
        token = dyadic_tokens[0]
        # Verify no gen field exists
        assert not hasattr(token, 'gen')

    def test_generate_direct_produces_configs(self):
        """generate_direct() produces PEConfigs with Instruction IRAM."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=3,
            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_cfg = result.pe_configs[0]
        assert pe_cfg.pe_id == 0
        assert 3 in pe_cfg.iram

        inst = pe_cfg.iram[3]
        assert isinstance(inst, Instruction)
        assert inst.opcode == ArithOp.ADD

    def test_generate_tokens_ordering(self):
        """generate_tokens() ordering: SM init → IRAM → ALLOC → frame slots → seeds."""
        data_def = IRDataDef(
            name="@data",
            sm_id=0,
            cell_addr=5,
            value=99,
            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(2, 1),
        )
        seed_node = IRNode(
            name="&const",
            opcode=RoutingOp.CONST,
            pe=0,
            iram_offset=10,
            act_id=0,
            const=42,
            mode=(OutputStyle.INHERIT, True, 1),
            wide=False,
            fref=0,
            seed=True,
            loc=SourceLoc(3, 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 seeds
        setup_count = sum(1 for t in tokens
                         if isinstance(t, (SMToken, PELocalWriteToken, FrameControlToken)))
        seed_count = sum(1 for t in tokens if isinstance(t, (DyadToken, MonadToken)))

        # Find indices
        if setup_count > 0 and seed_count > 0:
            first_setup = next(i for i, t in enumerate(tokens)
                              if isinstance(t, (SMToken, PELocalWriteToken, FrameControlToken)))
            first_seed = next(i for i, t in enumerate(tokens)
                             if isinstance(t, (DyadToken, MonadToken)))
            assert first_setup < first_seed


class TestAC65T0BootstrapPacking:
    """AC6.5: T0 bootstrap data uses pack_token() for packed flits."""

    def test_pack_token_for_t0_bootstrap(self):
        """Verify pack_token() can pack tokens for T0 bootstrap storage.

        In T0 bootstrap, tokens are written to T0 storage as packed flit sequences.
        Each flit is written via SMToken(WRITE) to T0 addresses.
        """
        # Create a sample dyadic token that might be bootstrapped to T0
        token = DyadToken(
            target=1,
            offset=10,
            act_id=2,
            data=0x1234,
            port=Port.L,
        )

        # Pack the token
        flits = pack_token(token)

        # Verify we get a sequence of flits
        assert isinstance(flits, list)
        assert len(flits) >= 1  # At least flit 1 (header)
        assert all(isinstance(f, int) for f in flits)

        # For T0 bootstrap, each flit would be written as:
        # SMToken(target=sm_id, addr=t0_offset + i, op=MemOp.WRITE, data=flit)
        for i, flit in enumerate(flits):
            # Verify flit is 16-bit
            assert 0 <= flit <= 0xFFFF, f"Flit {i} out of range: {flit}"

    def test_pack_token_monad_for_t0(self):
        """MonadToken can also be packed for T0 bootstrap."""
        token = MonadToken(
            target=2,
            offset=5,
            act_id=1,
            data=0x5678,
            inline=False,
        )

        flits = pack_token(token)
        assert len(flits) >= 1
        assert all(0 <= f <= 0xFFFF for f in flits)


class TestIntegration:
    """Integration tests for complete codegen pipeline."""

    def test_multinode_multiactivation(self):
        """Complex graph with multiple nodes and activations."""
        nodes = [
            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),
            ),
            IRNode(
                name="&b",
                opcode=ArithOp.SUB,
                pe=0,
                iram_offset=1,
                act_id=0,
                mode=(OutputStyle.INHERIT, False, 2),
                wide=False,
                fref=0,
                loc=SourceLoc(2, 1),
            ),
            IRNode(
                name="&c",
                opcode=ArithOp.INC,
                pe=1,
                iram_offset=0,
                act_id=1,
                mode=(OutputStyle.INHERIT, False, 1),
                wide=False,
                fref=0,
                loc=SourceLoc(3, 1),
            ),
        ]
        edges = [
            IREdge(source="&a", dest="&b", port=Port.L),
            IREdge(source="&b", dest="&c", port=Port.L),
        ]
        system = SystemConfig(pe_count=2, sm_count=1)
        graph = IRGraph(
            {node.name: node for node in nodes},
            edges=edges,
            system=system
        )

        result = generate_direct(graph)

        # Check pe_configs
        assert len(result.pe_configs) == 2
        assert result.pe_configs[0].pe_id == 0
        assert result.pe_configs[1].pe_id == 1

        # Check setup_tokens
        assert len(result.setup_tokens) > 0

        # Check that ALLOC tokens are generated for act_id=1
        alloc_tokens = [t for t in result.setup_tokens if isinstance(t, FrameControlToken)]
        assert any(t.act_id == 1 for t in alloc_tokens)


class TestAC6_4SetupTokensWithFrameDest:
    """AC6.4: Verify destination frame slot writes use pack_flit1() with is_dest=True."""

    def test_setup_tokens_frame_writes_use_pack_flit1(self):
        """Verify setup_tokens use pack_flit1() for frame destination writes.

        This test verifies that the codegen correctly uses pack_flit1() to encode
        FrameDest objects when writing to frame slots (region=1, is_dest=True).
        """
        from dataclasses import replace
        from asm.allocate import allocate
        from asm.ir import ResolvedDest

        # Create a simple graph with a node that has dest_l/dest_r
        source_node = IRNode(
            name="&source",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            act_id=0,
            mode=(OutputStyle.INHERIT, False, 1),
            wide=False,
            fref=0,
            loc=SourceLoc(1, 1),
        )

        dest_node = IRNode(
            name="&dest",
            opcode=ArithOp.SUB,
            pe=0,
            iram_offset=1,
            act_id=0,
            mode=(OutputStyle.SINK, False, 0),
            wide=False,
            fref=0,
            loc=SourceLoc(2, 1),
        )

        edge = IREdge(source="&source", dest="&dest", port=Port.L)

        system = SystemConfig(pe_count=1, sm_count=0)
        graph = IRGraph(
            nodes={"&source": source_node, "&dest": dest_node},
            edges=[edge],
            regions=[],
            system=system,
            call_sites=[],
        )

        # Run allocate to set up frame layouts and dest_l/dest_r
        allocated_graph = allocate(graph)
        if allocated_graph.errors:
            raise ValueError(f"Allocation errors: {allocated_graph.errors}")

        # Now run generate_direct on the allocated graph
        codegen_result = generate_direct(allocated_graph)

        # Find PELocalWriteToken entries with is_dest=True (frame dest writes)
        dest_write_tokens = [
            t for t in codegen_result.setup_tokens
            if isinstance(t, PELocalWriteToken) and t.is_dest
        ]

        # Should have at least one frame dest write
        assert len(dest_write_tokens) > 0, \
            "Should have PELocalWriteToken with is_dest=True for frame destinations"

        # Verify each dest write token has valid packed flit1 data
        for token in dest_write_tokens:
            assert token.region == 1, \
                f"Frame dest writes should use region=1 (frame), got {token.region}"
            # Data should be a valid packed flit1 (16-bit)
            assert 0 <= token.data <= 0xFFFF, \
                f"Frame dest data should be 16-bit, got {token.data}"