"""Tests for frame-based allocation in the assembler.

Tests AC5 acceptance criteria for frame layouts, IRAM deduplication,
activation ID assignment, and mode computation.
"""

from dataclasses import replace
from asm.allocate import (
    _assign_iram_offsets,
    _assign_act_ids,
    _deduplicate_iram,
    _compute_modes,
    _compute_frame_layouts,
    allocate,
)
from asm.ir import (
    IRNode,
    IREdge,
    SourceLoc,
    CallSite,
    IRGraph,
    SystemConfig,
    FrameLayout,
    FrameSlotMap,
)
from asm.errors import ErrorCategory, ErrorSeverity
from cm_inst import (
    ArithOp,
    MemOp,
    Port,
    RoutingOp,
    OutputStyle,
    TokenKind,
)


class TestIRAMDeduplication:
    """AC5.1: IRAM offsets deduplicated for identical instruction templates."""

    def test_identical_instructions_share_offset(self):
        """Two nodes with same opcode, mode, and fref share IRAM offset."""
        # Create two nodes with identical template
        node1 = IRNode(
            name="&add1",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, False, 1),
            fref=0,
            wide=False,
        )
        node2 = IRNode(
            name="&add2",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=1,
            mode=(OutputStyle.INHERIT, False, 1),
            fref=0,
            wide=False,
        )

        nodes_on_pe = {"&add1": node1, "&add2": node2}
        deduped = _deduplicate_iram(nodes_on_pe, pe_id=0)

        # Both should have same offset
        assert deduped["&add1"].iram_offset == 0
        assert deduped["&add2"].iram_offset == 0

    def test_different_opcodes_different_offsets(self):
        """Nodes with different opcodes keep different offsets."""
        node1 = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, False, 1),
            fref=0,
            wide=False,
        )
        node2 = IRNode(
            name="&sub",
            opcode=ArithOp.SUB,
            pe=0,
            iram_offset=1,
            mode=(OutputStyle.INHERIT, False, 1),
            fref=0,
            wide=False,
        )

        nodes_on_pe = {"&add": node1, "&sub": node2}
        deduped = _deduplicate_iram(nodes_on_pe, pe_id=0)

        assert deduped["&add"].iram_offset == 0
        assert deduped["&sub"].iram_offset == 1

    def test_same_opcode_different_modes_different_offsets(self):
        """Same opcode but different modes get different offsets."""
        node1 = IRNode(
            name="&add1",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, False, 1),
            fref=0,
            wide=False,
        )
        node2 = IRNode(
            name="&add2",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=1,
            mode=(OutputStyle.INHERIT, True, 1),  # has_const differs
            fref=0,
            wide=False,
        )

        nodes_on_pe = {"&add1": node1, "&add2": node2}
        deduped = _deduplicate_iram(nodes_on_pe, pe_id=0)

        assert deduped["&add1"].iram_offset == 0
        assert deduped["&add2"].iram_offset == 1

    def test_seed_nodes_excluded(self):
        """Seed nodes are excluded from deduplication."""
        seed_node = IRNode(
            name="&seed",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=None,
            seed=True,
        )
        normal_node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, False, 1),
            fref=0,
            wide=False,
        )

        nodes_on_pe = {"&seed": seed_node, "&add": normal_node}
        deduped = _deduplicate_iram(nodes_on_pe, pe_id=0)

        # Seed node should be unchanged
        assert deduped["&seed"].seed
        assert deduped["&seed"].iram_offset is None


class TestActivationIDAssignment:
    """AC5.3, AC5.7: Activation IDs assigned sequentially, exhaustion checked."""

    def test_single_function_gets_act_id_0(self):
        """Single function scope gets act_id=0."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
        )

        updated, errors = _assign_act_ids([node], {"&add": node}, frame_count=8, pe_id=0)

        # Root scope node should get act_id=0
        assert updated["&add"].act_id == 0
        assert not errors

    def test_multiple_functions_different_act_ids(self):
        """Different functions on same PE get different act_ids."""
        node1 = IRNode(
            name="$func1.&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
        )
        node2 = IRNode(
            name="$func2.&sub",
            opcode=ArithOp.SUB,
            pe=0,
            iram_offset=1,
        )

        all_nodes = {"$func1.&add": node1, "$func2.&sub": node2}
        updated, errors = _assign_act_ids(
            [node1, node2], all_nodes, frame_count=8, pe_id=0
        )

        # Different functions should get different act_ids
        assert updated["$func1.&add"].act_id == 0
        assert updated["$func2.&sub"].act_id == 1
        assert not errors

    def test_activation_exhaustion_error(self):
        """When >frame_count act_ids needed, error with FRAME category."""
        nodes = []
        all_nodes = {}
        for i in range(10):
            node = IRNode(
                name=f"$func{i}.&op",
                opcode=ArithOp.ADD,
                pe=0,
                iram_offset=i,
            )
            nodes.append(node)
            all_nodes[f"$func{i}.&op"] = node

        updated, errors = _assign_act_ids(nodes, all_nodes, frame_count=8, pe_id=0)

        # Should report activation exhaustion error
        assert len(errors) > 0
        assert errors[0].category == ErrorCategory.FRAME
        assert "exhaustion" in errors[0].message.lower()

    def test_same_function_across_pes_reuses_act_id(self):
        """Same function on different PEs can reuse act_id."""
        node1 = IRNode(
            name="$func.&add",
            opcode=ArithOp.ADD,
            pe=0,
            iram_offset=0,
        )
        node2 = IRNode(
            name="$func.&sub",
            opcode=ArithOp.SUB,
            pe=1,
            iram_offset=0,
        )

        # Process PE 0
        all_nodes = {"$func.&add": node1, "$func.&sub": node2}
        updated0, errors0 = _assign_act_ids([node1], all_nodes, frame_count=8, pe_id=0)

        # Process PE 1
        updated1, errors1 = _assign_act_ids([node2], all_nodes, frame_count=8, pe_id=1)

        # Both can have same act_id within their respective PEs
        assert updated0["$func.&add"].act_id == 0
        assert updated1["$func.&sub"].act_id == 0


class TestModeComputation:
    """AC5.5: Mode (OutputStyle + has_const + dest_count) computed from topology."""

    def test_sink_op_has_zero_dest_count(self):
        """SM WRITE produces SINK mode with dest_count=0."""
        node = IRNode(
            name="&write",
            opcode=MemOp.WRITE,
            const=0x100,
            pe=0,
        )

        edges_by_source = {"&write": []}  # No outgoing edges
        modes = _compute_modes({"&write": node}, edges_by_source)

        mode = modes["&write"].mode
        assert mode[0] == OutputStyle.SINK  # OutputStyle
        assert mode[2] == 0  # dest_count should be 0

    def test_read_op_has_inherit_style(self):
        """SM READ produces INHERIT style."""
        node = IRNode(
            name="&read",
            opcode=MemOp.READ,
            pe=0,
        )

        edges_by_source = {"&read": []}
        modes = _compute_modes({"&read": node}, edges_by_source)

        mode = modes["&read"].mode
        assert mode[0] == OutputStyle.INHERIT

    def test_has_const_flag_set(self):
        """Node with const field has has_const=True."""
        node = IRNode(
            name="&add_const",
            opcode=ArithOp.ADD,
            const=42,
            pe=0,
        )

        edges_by_source = {"&add_const": []}
        modes = _compute_modes({"&add_const": node}, edges_by_source)

        mode = modes["&add_const"].mode
        assert mode[1] is True  # has_const

    def test_dest_count_from_edges(self):
        """dest_count matches number of outgoing edges."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
        )
        edge1 = IREdge(source="&add", dest="&sub", port=Port.L)
        edge2 = IREdge(source="&add", dest="&mul", port=Port.R)

        edges_by_source = {"&add": [edge1, edge2]}
        modes = _compute_modes({"&add": node}, edges_by_source)

        mode = modes["&add"].mode
        assert mode[2] == 2  # dest_count

    def test_free_frame_produces_sink(self):
        """FREE_FRAME produces SINK style."""
        node = IRNode(
            name="&free",
            opcode=RoutingOp.FREE_FRAME,
            pe=0,
        )

        edges_by_source = {"&free": []}
        modes = _compute_modes({"&free": node}, edges_by_source)

        mode = modes["&free"].mode
        assert mode[0] == OutputStyle.SINK

    def test_ctx_override_produces_change_tag_with_dest_count_1(self):
        """AC5.5: ctx_override=True on outgoing edge produces CHANGE_TAG mode with dest_count=1.

        When a node has an outgoing edge with ctx_override=True (crossing function boundary),
        the output mode should be CHANGE_TAG with dest_count=1.
        """
        source_node = IRNode(
            name="&source",
            opcode=ArithOp.ADD,
            pe=0,
        )
        dest_node = IRNode(
            name="&dest",
            opcode=ArithOp.SUB,
            pe=0,
        )

        # Create edge with ctx_override=True
        edge = IREdge(
            source="&source",
            dest="&dest",
            port=Port.L,
            ctx_override=True,
        )

        edges_by_source = {"&source": [edge], "&dest": []}
        nodes = {"&source": source_node, "&dest": dest_node}
        modes = _compute_modes(nodes, edges_by_source)

        mode = modes["&source"].mode
        # ctx_override=True should produce CHANGE_TAG style
        assert mode[0] == OutputStyle.CHANGE_TAG, \
            f"ctx_override=True should produce CHANGE_TAG, got {mode[0]}"
        # dest_count should be 1 (one outgoing edge)
        assert mode[2] == 1, \
            f"dest_count should be 1 for single ctx_override edge, got {mode[2]}"


class TestFrameLayoutComputation:
    """AC5.2, AC5.4, AC5.6: Frame slot assignment and deduplication."""

    def test_frame_layout_canonical_per_function(self):
        """Two activations of same function produce identical FrameLayout."""
        # Create nodes with same structure but different names
        node1 = IRNode(
            name="&add1",
            opcode=ArithOp.ADD,
            pe=0,
            act_id=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, False, 1),
        )
        node2 = IRNode(
            name="&add2",
            opcode=ArithOp.ADD,
            pe=0,
            act_id=1,
            iram_offset=1,
            mode=(OutputStyle.INHERIT, False, 1),
        )

        nodes_on_pe = {"&add1": node1, "&add2": node2}
        all_nodes = nodes_on_pe.copy()
        edges_by_source = {"&add1": [], "&add2": []}
        edges_by_dest = {}

        updated, errors = _compute_frame_layouts(
            nodes_on_pe,
            edges_by_source,
            edges_by_dest,
            all_nodes,
            frame_slots=64,
            matchable_offsets=8,
            pe_id=0,
        )

        # Both activations should have layouts with same structure
        layout1 = updated["&add1"].frame_layout
        layout2 = updated["&add2"].frame_layout
        assert layout1 is not None
        assert layout2 is not None
        # Layouts should have same slot structure (though may be different objects)
        assert layout1.total_slots == layout2.total_slots

    def test_frame_slot_overflow_error(self):
        """When total slots > frame_slots, error with FRAME category."""
        # Create nodes that would exceed frame slots
        nodes_on_pe = {}
        all_nodes = {}
        for i in range(100):
            node = IRNode(
                name=f"&op{i}",
                opcode=ArithOp.ADD,
                pe=0,
                act_id=0,
                iram_offset=i,
                mode=(OutputStyle.INHERIT, True, 2),  # Each needs 3+ slots
            )
            nodes_on_pe[f"&op{i}"] = node
            all_nodes[f"&op{i}"] = node

        edges_by_source = {f"&op{i}": [] for i in range(100)}
        edges_by_dest = {}

        updated, errors = _compute_frame_layouts(
            nodes_on_pe,
            edges_by_source,
            edges_by_dest,
            all_nodes,
            frame_slots=64,
            matchable_offsets=8,
            pe_id=0,
        )

        # Should report frame slot overflow
        assert len(errors) > 0
        assert errors[0].category == ErrorCategory.FRAME

    def test_match_slots_at_low_offsets(self):
        """Match operands occupy slots 0 to matchable_offsets-1."""
        # Create dyadic (matchable) node
        node1 = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            act_id=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, False, 1),
        )
        # Monadic node shouldn't count as match
        node2 = IRNode(
            name="&inc",
            opcode=ArithOp.INC,
            pe=0,
            act_id=0,
            iram_offset=1,
            mode=(OutputStyle.INHERIT, False, 1),
        )

        nodes_on_pe = {"&add": node1, "&inc": node2}
        all_nodes = nodes_on_pe.copy()
        edges_by_source = {"&add": [], "&inc": []}
        edges_by_dest = {}

        updated, errors = _compute_frame_layouts(
            nodes_on_pe,
            edges_by_source,
            edges_by_dest,
            all_nodes,
            frame_slots=64,
            matchable_offsets=8,
            pe_id=0,
        )

        assert not errors
        layout = updated["&add"].frame_layout
        assert layout is not None
        # Match slots should start at 0
        assert 0 in layout.slot_map.match_slots

    def test_matchable_offset_exceedance_warning(self):
        """AC5.8: Warn when dyadic_count exceeds matchable_offsets.

        Creates an activation with 9+ dyadic nodes but matchable_offsets=8.
        Verifies: (1) errors list contains warning with FRAME category and WARNING severity,
        (2) the function still returns a valid layout (not a hard error).
        """
        # Create 9 dyadic nodes (more than matchable_offsets=8)
        nodes_on_pe = {}
        all_nodes = {}
        for i in range(9):
            node = IRNode(
                name=f"&add{i}",
                opcode=ArithOp.ADD,  # Dyadic
                pe=0,
                act_id=0,
                iram_offset=i,
                mode=(OutputStyle.INHERIT, False, 1),
            )
            nodes_on_pe[f"&add{i}"] = node
            all_nodes[f"&add{i}"] = node

        edges_by_source = {f"&add{i}": [] for i in range(9)}
        edges_by_dest = {}

        updated, errors = _compute_frame_layouts(
            nodes_on_pe,
            edges_by_source,
            edges_by_dest,
            all_nodes,
            frame_slots=64,
            matchable_offsets=8,  # Less than 9 dyadic nodes
            pe_id=0,
        )

        # Should generate a warning with FRAME category and WARNING severity
        assert len(errors) > 0, "Expected warning about matchable offset exceedance"
        warning = errors[0]
        assert warning.category == ErrorCategory.FRAME
        assert warning.severity == ErrorSeverity.WARNING
        assert "dyadic" in warning.message.lower()
        assert "match slot" in warning.message.lower()

        # Function should still return valid layout (not a hard error)
        assert updated is not None
        assert len(updated) == 9
        for i in range(9):
            assert f"&add{i}" in updated
            assert updated[f"&add{i}"].frame_layout is not None

    def test_constants_deduplicated_by_value(self):
        """Instructions with same constant value can share slot."""
        node1 = IRNode(
            name="&op1",
            opcode=ArithOp.ADD,
            const=42,
            pe=0,
            act_id=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, True, 1),
        )
        node2 = IRNode(
            name="&op2",
            opcode=ArithOp.SUB,
            const=42,  # Same constant
            pe=0,
            act_id=0,
            iram_offset=1,
            mode=(OutputStyle.INHERIT, True, 1),
        )

        nodes_on_pe = {"&op1": node1, "&op2": node2}
        all_nodes = nodes_on_pe.copy()
        edges_by_source = {"&op1": [], "&op2": []}
        edges_by_dest = {}

        updated, errors = _compute_frame_layouts(
            nodes_on_pe,
            edges_by_source,
            edges_by_dest,
            all_nodes,
            frame_slots=64,
            matchable_offsets=8,
            pe_id=0,
        )

        assert not errors
        # Both should have same layout with one const slot
        layout = updated["&op1"].frame_layout
        assert len(layout.slot_map.const_slots) == 1

    def test_per_node_fref_differentiation(self):
        """AC5.2: Nodes in same activation with different modes get different frefs.

        Verifies that each node gets a unique fref based on its mode's slot count.
        Node A: has_const=True, dest_count=1 → 2 slots → fref=8
        Node B: has_const=False, dest_count=1 → 1 slot  → fref=10 (8+2)
        """
        # Create two nodes with different modes
        node_a = IRNode(
            name="&op_a",
            opcode=ArithOp.ADD,
            const=42,  # has_const=True
            pe=0,
            act_id=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, True, 1),  # has_const=True, 1 dest → 2 slots
        )
        node_b = IRNode(
            name="&op_b",
            opcode=ArithOp.SUB,
            pe=0,
            act_id=0,
            iram_offset=1,
            mode=(OutputStyle.INHERIT, False, 1),  # has_const=False, 1 dest → 1 slot
        )

        # Create edge from &op_a to &op_b
        edge = IREdge(source="&op_a", dest="&op_b", port=Port.L)

        nodes_on_pe = {"&op_a": node_a, "&op_b": node_b}
        all_nodes = nodes_on_pe.copy()
        edges_by_source = {"&op_a": [edge], "&op_b": []}
        edges_by_dest = {"&op_b": [edge]}

        matchable_offsets = 8
        updated, errors = _compute_frame_layouts(
            nodes_on_pe,
            edges_by_source,
            edges_by_dest,
            all_nodes,
            frame_slots=64,
            matchable_offsets=matchable_offsets,
            pe_id=0,
        )

        assert not errors, f"Unexpected errors: {errors}"

        # Verify both nodes got frefs
        assert updated["&op_a"].fref is not None
        assert updated["&op_b"].fref is not None

        # Verify they have different frefs
        fref_a = updated["&op_a"].fref
        fref_b = updated["&op_b"].fref
        assert fref_a != fref_b, f"Expected different frefs, got {fref_a} and {fref_b}"

        # Verify fref_a starts at matchable_offsets
        assert fref_a == matchable_offsets, f"Expected fref_a={matchable_offsets}, got {fref_a}"

        # Verify fref_b = fref_a + slot_count_a (where slot_count_a = 2)
        # Mode A: (INHERIT, True, 1) → 1 (const) + 1 (dest) = 2 slots
        expected_fref_b = fref_a + 2
        assert fref_b == expected_fref_b, f"Expected fref_b={expected_fref_b}, got {fref_b}"


class TestFrameLayoutIntegration:
    """Integration tests for full allocation pipeline."""

    def test_minimal_graph(self):
        """Minimal graph with one node allocates successfully."""
        node = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            pe=0,
            act_id=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, False, 1),
        )

        updated, errors = _compute_frame_layouts(
            {"&add": node},
            {"&add": []},
            {},
            {"&add": node},
            frame_slots=64,
            matchable_offsets=8,
            pe_id=0,
        )

        assert not errors
        assert updated["&add"].frame_layout is not None
        assert updated["&add"].fref is not None

    def test_multiple_activations_separate_layouts(self):
        """Nodes with different act_ids get separate layout configurations."""
        node1 = IRNode(
            name="&op1",
            opcode=ArithOp.ADD,
            pe=0,
            act_id=0,
            iram_offset=0,
            mode=(OutputStyle.INHERIT, False, 1),
        )
        node2 = IRNode(
            name="&op2",
            opcode=ArithOp.ADD,
            pe=0,
            act_id=1,  # Different activation
            iram_offset=1,
            mode=(OutputStyle.INHERIT, True, 1),  # Different mode
        )

        updated, errors = _compute_frame_layouts(
            {"&op1": node1, "&op2": node2},
            {"&op1": [], "&op2": []},
            {},
            {"&op1": node1, "&op2": node2},
            frame_slots=64,
            matchable_offsets=8,
            pe_id=0,
        )

        assert not errors
        layout1 = updated["&op1"].frame_layout
        layout2 = updated["&op2"].frame_layout
        # Layouts can have different sizes depending on mode
        # (one has const, one doesn't)
        assert layout1 is not None
        assert layout2 is not None

    def test_full_allocate_pipeline(self):
        """Full integration test: allocate() with complete IRGraph.

        Tests that allocate() correctly:
        - Assigns IRAM offsets
        - Assigns activation IDs
        - Computes modes
        - Computes frame layouts
        - Deduplicates IRAM entries
        - Resolves destinations to FrameDest
        """
        # Create a simple graph: two arithmetic nodes with an edge
        node_add = IRNode(
            name="&add",
            opcode=ArithOp.ADD,
            const=None,
            pe=0,
            loc=SourceLoc(1, 1),
        )
        node_inc = IRNode(
            name="&inc",
            opcode=ArithOp.INC,
            const=None,
            pe=0,
            loc=SourceLoc(2, 1),
        )
        edge = IREdge(
            source="&add",
            dest="&inc",
            port=Port.L,
            loc=SourceLoc(1, 5),
        )

        # Create the IRGraph
        system = SystemConfig(
            pe_count=1,
            sm_count=0,
            iram_capacity=64,
            frame_count=8,
            frame_slots=64,
            matchable_offsets=8,
        )

        graph = IRGraph(
            nodes={"&add": node_add, "&inc": node_inc},
            edges=[edge],
            regions=[],
            system=system,
            call_sites=[],
        )

        # Call allocate()
        result = allocate(graph)

        # Verify no errors
        assert not result.errors, f"Unexpected errors: {result.errors}"

        add_node = result.nodes["&add"]
        inc_node = result.nodes["&inc"]

        # IRAM offsets: &add is dyadic -> offset 0, &inc is monadic -> offset 1
        assert add_node.iram_offset == 0
        assert inc_node.iram_offset == 1

        # Both in root scope -> act_id 0
        assert add_node.act_id == 0
        assert inc_node.act_id == 0

        # &add has 1 outgoing edge, no const -> INHERIT, False, 1
        assert add_node.mode == (OutputStyle.INHERIT, False, 1)
        # &inc has 0 outgoing edges, no const -> SINK, False, 0
        assert inc_node.mode == (OutputStyle.SINK, False, 0)

        # Frame layouts assigned
        assert add_node.frame_layout is not None
        assert inc_node.frame_layout is not None

        # Frefs: both in same activation, &add needs 1 dest slot, &inc needs 0
        assert add_node.fref is not None
        assert inc_node.fref is not None

        # &add dest_l resolves to &inc
        assert add_node.dest_l is not None
        assert add_node.dest_l.frame_dest is not None
        assert add_node.dest_l.frame_dest.target_pe == 0
        assert add_node.dest_l.frame_dest.offset == 1
        assert add_node.dest_l.frame_dest.act_id == 0
        assert add_node.dest_l.frame_dest.port == Port.L
        assert add_node.dest_l.frame_dest.token_kind == TokenKind.MONADIC