"""Phase 1 spec tests for ADR-0065 P2 — `softmax_merge` recipe lowering.

P2 adds `tl_recipes.py` (RECIPE_DESCRIPTORS) and a TLContext lowering pass
so that

    tl.composite(
        op="gemm", b=V_ref, out=O,
        prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
        epilogue=[{"op": "add", "other": O}],
    )

expands the single `softmax_merge` recipe into 8 flat MATH OpSpecs, binds the
recipe's primary output `P` into the head GEMM's `a` operand, and emits one
`CompositeCmd` with 10 ops + `rw_handles == (m, l, O)` (ADR-0065 D5 / DDD §6).

This is host-side lowering only — PE_SCHEDULER consumption of the pre-GEMM
KERNEL-scope ops is P3, so these tests inspect the lowered CompositeCmd
object and do NOT run it through the scheduler.

Phase 1 (this commit): tests only. FAIL until P2:
  - `kernbench.triton_emu.tl_recipes` does not exist,
  - `TLContext.composite` has no `prologue=` / `out=TensorHandle` kwargs.
"""
from __future__ import annotations

import math

from kernbench.common.pe_commands import CompositeCmd, TensorHandle
from kernbench.triton_emu.tl_context import TLContext

G, TILE, D = 8, 64, 128

_EXPECTED_KINDS = [
    "rmax", "max_elem", "exp_diff", "exp_diff", "rsum", "fma",
    "mul_bcast", "copy",            # 8 MATH (softmax_merge)
    "gemm",                          # head GEMM (P · V)
    "add",                           # epilogue
]


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


def _lower():
    """Lower decode-opt2 #2 and return (tl, cmd, handles)."""
    tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000,
                   scratch_size=1 << 20)
    s = _tcm(0x1000, (G, TILE))      # Sj — #1's output, TCM-resident
    m = _tcm(0x2000, (G,))
    l = _tcm(0x3000, (G,))
    O = _tcm(0x4000, (G, D))
    V = tl.ref(0x5000, shape=(TILE, D), dtype="f16")
    tl.composite(
        op="gemm", b=V, out=O,
        prologue=[{"op": "softmax_merge", "s": s, "m": m, "l": l, "O": O}],
        epilogue=[{"op": "add", "other": O}],
    )
    cmd = [c for c in tl.commands if isinstance(c, CompositeCmd)][-1]
    return tl, cmd, dict(s=s, m=m, l=l, O=O, V=V)


# ── recipe shape (ADR-0065 Test #2) ──────────────────────────────────


def test_recipe_lowers_to_ten_ops_in_order():
    _, cmd, _ = _lower()
    assert [o.kind for o in cmd.ops] == _EXPECTED_KINDS, [o.kind for o in cmd.ops]


def test_recipe_rw_handles_are_m_l_O():
    _, cmd, h = _lower()
    assert cmd.rw_handles == (h["m"], h["l"], h["O"]), cmd.rw_handles


def test_head_gemm_autobinds_primary_out_and_geometry():
    _, cmd, h = _lower()
    gemm = cmd.ops[8]
    assert gemm.kind == "gemm"
    # b is V (the kernel-provided operand); a is the recipe's primary out P
    assert gemm.operands["b"] == h["V"]
    assert "a" in gemm.operands and gemm.operands["a"].shape == (G, TILE)
    assert gemm.out is not None and gemm.out.addr == h["O"].addr
    assert gemm.extra.get("m") == G
    assert gemm.extra.get("k") == TILE
    assert gemm.extra.get("n") == D


def test_recipe_scratch_addresses_are_concrete_and_distinct():
    _, cmd, _ = _lower()
    # Collect every operand/out handle address across the 8 MATH ops.
    addrs = set()
    for op in cmd.ops[:8]:
        for hd in op.operands.values():
            assert hd.addr is not None
            addrs.add(hd.addr)
        if op.out is not None:
            addrs.add(op.out.addr)
    # The 5 scratch slots (m_loc, m_new, corr, P, l_loc) must have been
    # allocated to distinct, non-zero addresses (scratch_base supplied).
    scratch_addrs = {a for a in addrs if a >= 0x100000}
    assert len(scratch_addrs) >= 5, scratch_addrs


def test_recipe_composite_within_size_cap():
    """ADR-0065 Test #9 — the 10-op composite fits the default 1024 cap."""
    _, cmd, _ = _lower()
    assert cmd.logical_bytes < 1024, cmd.logical_bytes