"""Phase 1 spec tests for ADR-0065 P3 — flat-ops PE_SCHEDULER plan.

P3 rewrites tile-plan generation to scan the flat ops list for the GEMM
(≤1) and place pre-GEMM KERNEL ops as single-shot prologue MATH stages
(ADR-0065 D3). A new tiling entry ``generate_plan_from_ops(ops, ...)``
produces a ``PipelinePlan`` with ``prologue_stages`` / ``epilogue_stages``.

DMA insertion keeps the existing **``pinned``** signal (already-in-TCM →
skip DMA; not-pinned → stream from HBM), NOT a ``space`` flip — so existing
bench Stage sequences stay byte-equal. Recipe TCM scratch / primary-out
handles are marked ``pinned=True`` so the head GEMM's auto-bound ``a`` (= the
recipe's ``P``) is consumed in place, while ``b`` (a ``tl.ref``) is streamed.

These tests target the tiling entry directly (no engine), building OpSpecs
by hand so each operand's ``pinned`` flag is controlled precisely.

Phase 1 (this commit): tests only. FAIL until P3:
  - ``generate_plan_from_ops`` does not exist,
  - ``PipelinePlan`` has no ``prologue_stages`` field.
"""
from __future__ import annotations

import math

from kernbench.common.pe_commands import OpSpec, Scope, TensorHandle

PE = "sip0.cube0.pe0"
TM = TK = TN = 64


def _h(addr: int, shape: tuple[int, ...], pinned: bool = False) -> TensorHandle:
    return TensorHandle(
        id=f"h{addr:x}", addr=addr, shape=shape, dtype="f16",
        nbytes=2 * math.prod(shape), space="tcm", pinned=pinned,
    )


def _gemm_op(a, b, out) -> OpSpec:
    return OpSpec(
        kind="gemm", scope=Scope.OUTPUT_TILE,
        operands={"a": a, "b": b}, out=out,
        extra={"m": a.shape[-2], "k": a.shape[-1], "n": b.shape[-1]},
    )


def _plan(ops):
    from kernbench.components.builtin.tiling import generate_plan_from_ops

    return generate_plan_from_ops(
        ops=tuple(ops), tile_m=TM, tile_k=TK, tile_n=TN,
        bytes_per_element=2, pe_prefix=PE,
    )


def _stage_types(stages):
    return [s.stage_type.name for s in stages]


# ── DMA from the pinned flag (byte-equal with generate_gemm_plan) ─────


def test_unpinned_operand_gets_dma_read():
    from kernbench.components.builtin.pe_types import StageType

    a = _h(0x1000, (TM, TK), pinned=True)            # already in TCM
    b = _h(0x2000, (TK, TN), pinned=False)           # streamed from HBM
    out = _h(0x3000, (TM, TN))
    tile = _plan([_gemm_op(a, b, out)]).tiles[0]
    reads = [s for s in tile.stages if s.stage_type == StageType.DMA_READ]
    assert len(reads) == 1, _stage_types(tile.stages)   # only B


def test_pinned_operands_skip_dma():
    from kernbench.components.builtin.pe_types import StageType

    a = _h(0x1000, (TM, TK), pinned=True)
    b = _h(0x2000, (TK, TN), pinned=True)
    out = _h(0x3000, (TM, TN))
    tile = _plan([_gemm_op(a, b, out)]).tiles[0]
    assert not [s for s in tile.stages if s.stage_type == StageType.DMA_READ]


# ── prologue placement (ADR-0065 D3) ─────────────────────────────────


def test_recipe_prologue_math_ops_become_prologue_stages():
    from kernbench.components.builtin.pe_types import StageType

    # 8 KERNEL-scope MATH ops, then the head GEMM whose `a` is the recipe's
    # primary-out P (pinned → in place) and `b` is a streamed ref.
    sc = _h(0x100000, (8, 64), pinned=True)
    pre = [
        OpSpec(kind=k, scope=Scope.KERNEL, operands={"src": sc}, out=sc)
        for k in ("rmax", "max_elem", "exp_diff", "exp_diff",
                  "rsum", "fma", "mul_bcast", "copy")
    ]
    P = _h(0x100400, (8, 64), pinned=True)           # primary-out, in TCM
    V = _h(0x2000, (64, 128), pinned=False)          # streamed
    out = _h(0x3000, (8, 128))
    plan = _plan([*pre, _gemm_op(P, V, out)])

    # 8 single-shot MATH prologue stages, none of them DMA.
    assert len(plan.prologue_stages) == 8, _stage_types(plan.prologue_stages)
    assert all(s.stage_type == StageType.MATH for s in plan.prologue_stages)
    assert not any(s.stage_type in (StageType.DMA_READ, StageType.DMA_WRITE)
                   for s in plan.prologue_stages)
    # Head GEMM drives the tile loop; only V (not pinned) gets a DMA_READ.
    reads = [s for s in plan.tiles[0].stages
             if s.stage_type == StageType.DMA_READ]
    assert len(reads) == 1, _stage_types(plan.tiles[0].stages)


def test_math_only_composite_has_no_gemm_tiles():
    """A legacy MATH composite (no GEMM op) → math tile plan, no GEMM stage.

    Shaped like ``tl.composite(op="math", a=..., math_op="exp")`` — head
    operand keyed ``"a"`` with the op kind in ``extra["math_op"]``."""
    from kernbench.components.builtin.pe_types import StageType

    src = _h(0x1000, (TM, TN), pinned=False)
    out = _h(0x3000, (TM, TN))
    op = OpSpec(kind="math", scope=Scope.OUTPUT_TILE,
                operands={"a": src}, extra={"math_op": "exp"}, out=out)
    plan = _plan([op])
    all_stages = []
    for t in plan.tiles:
        all_stages.extend(t.stages)
    assert plan.tiles, "math-only composite must still produce tiles"
    assert not any(s.stage_type == StageType.GEMM for s in all_stages)