"""TLContext: fake Triton Language module for kernel performance simulation.

Passed as the `tl` parameter to kernel functions. Each API call records a
PeCommand in the internal trace. After the kernel returns, PE_CPU replays
the command list through SimPy.

Kernel code looks like standard Python — no yield, no async:

    def my_kernel(a_ptr, b_ptr, out_ptr, tl):
        pid = tl.program_id(0)
        a = tl.load(a_ptr, shape=(32, 64), dtype="f16")
        b = tl.load(b_ptr + pid * stride, shape=(64, 32), dtype="f16")
        tl.composite(op="gemm", a=a, b=b, out_ptr=out_ptr)
"""
from __future__ import annotations

import math
from typing import Literal

from kernbench.common.ipcq_types import IpcqRecvCmd, IpcqSendCmd, RecvFuture
from kernbench.common.pe_commands import (
    EPILOGUE_OPS,
    CompletionHandle,
    CompositeCmd,
    CopyCmd,
    DmaReadCmd,
    DmaWriteCmd,
    GemmCmd,
    MathCmd,
    OpSpec,
    PeCommand,
    PeCpuOverheadCmd,
    Scope,
    TensorHandle,
    WaitCmd,
)

_DTYPE_BYTES: dict[str, int] = {
    "f16": 2, "f32": 4, "f64": 8,
    "bf16": 2,
    "i8": 1, "i16": 2, "i32": 4, "i64": 8,
    "u8": 1, "u16": 2, "u32": 4, "u64": 8,
}


class LoadFuture:
    """Lazy ``tl.load`` future (ADR-0062 §D2).

    Mirrors ``RecvFuture`` for IPCQ, generalised to HBM loads. Carries
    the originating ``DmaReadCmd``, the SimPy completion event (set by
    the runner once the DMA has been issued), and a resolved flag.

    Consumer ops auto-wait via ``TLContext._await_pending(handle)``,
    which yields ``event`` if ``resolved`` is False, then reads the
    DMA's bytes into ``handle.data`` and marks the future resolved.
    """

    __slots__ = ("cmd", "event", "resolved", "data")

    def __init__(self, cmd: DmaReadCmd) -> None:
        self.cmd = cmd
        self.event: object | None = None  # simpy.Event set by runner
        self.resolved: bool = False
        self.data: object = None


class CompositeFuture:
    """Pending marker on a composite's output handle (ADR-0065 D8).

    Carries the composite's ``CompletionHandle`` so a downstream op (via
    ``_await_pending``) or ``tl.wait`` can await it. Distinct from
    ``LoadFuture``: composite completion is tracked by completion id in the
    runner's pending dict, so awaiting just re-issues a ``WaitCmd`` (a
    second wait on an already-drained completion is a harmless no-op).
    """

    __slots__ = ("completion", "resolved")

    def __init__(self, completion: "CompletionHandle") -> None:
        self.completion = completion
        self.resolved: bool = False


class _ScratchScope:
    """Context manager that recycles per-tile scratch (ADR-0063 D1).

    ``__enter__`` snapshots ``_scratch_cursor``; ``__exit__`` restores it,
    so every handle allocated inside the ``with``-block has its address
    freed for the next iteration. Persistent state (running ``(m, ℓ, O)``,
    prefetch buffers) lives outside the scope per ADR-0063 D3.
    """

    def __init__(self, ctx: "TLContext") -> None:
        self._ctx = ctx
        self._save: int | None = None

    def __enter__(self) -> "_ScratchScope":
        self._save = self._ctx._scratch_cursor
        return self

    def __exit__(self, *exc_info: object) -> bool:
        if self._save is not None:
            self._ctx._scratch_cursor = self._save
        return False


class TLContext:
    """Fake Triton Language context.

    Args:
        pe_id: program instance index (returned by program_id).
        num_programs: total number of program instances.
        dispatch_cycles: uniform PE_CPU overhead per dispatched command.
            Back-compat fallback used only when ``cost_model`` is None
            (ADR-0046 §D6).
        cost_model: structural dispatch cost model (ADR-0064 Rev2). When
            provided, every dispatched PeCommand is charged
            ``FIXED + cmd.logical_bytes × R`` cycles (÷ ``clock_freq_ghz``
            for ns) as a ``PeCpuOverheadCmd`` emitted just before it. Live
            PE_CPU paths construct TLContext with the per-PE model so the
            hybrid's CPU-saturation lever is measurable.
        clock_freq_ghz: cycle→ns conversion for the cost model (ADR-0064 D3).
    """

    def __init__(
        self,
        pe_id: int = 0,
        num_programs: int = 1,
        dispatch_cycles: int = 1,
        runner: Any = None,
        cube_id: int = 0,
        num_cubes: int = 1,
        scratch_base: int = 0,
        scratch_size: int = 1 << 20,  # 1 MiB per kernel invocation
        cost_model: "PeCostModel | None" = None,
        clock_freq_ghz: float = 1.0,
    ) -> None:
        self._pe_id = pe_id
        self._num_programs = num_programs
        self._cube_id = cube_id
        self._num_cubes = num_cubes
        self._dispatch_cycles = dispatch_cycles
        self._cost_model = cost_model
        self._clock_freq_ghz = clock_freq_ghz
        self._commands: list[PeCommand] = []
        self._handle_counter = 0
        self._completion_counter = 0
        self._runner = runner  # KernelRunner for greenlet mode (ADR-0020 D3)
        # PE-local scratch allocator for math/compute output handles.
        # Each binary/unary/reduction op auto-allocates a unique addr from
        # this pool so the resulting TensorHandle can be the source of a
        # later tl.send / tl.store. Cursor resets on every kernel invocation.
        self._scratch_base = scratch_base
        self._scratch_size = scratch_size
        self._scratch_cursor = 0

    def scratch_scope(self) -> _ScratchScope:
        """Per-tile scratch recycling context manager (ADR-0063).

        Usage:

            with tl.scratch_scope():
                s   = tl.dot(q, k_t)        # per-tile temporaries —
                p   = tl.softmax(s)         # their scratch is rewound
                o_j = tl.dot(p, v)          # on __exit__

        Persistent state (running ``(m, ℓ, O)``, lazy-load prefetch
        buffers) must be allocated **outside** the scope; only handles
        allocated inside are recycled.
        """
        return _ScratchScope(self)

    def _scratch_alloc(self, nbytes: int) -> int:
        """Allocate a unique scratch address for an output TensorHandle.

        Returns 0 if no scratch base was configured (e.g. command-list mode);
        in that case the resulting handle has addr=0 and cannot be used as a
        send/store source. Greenlet/runner mode always supplies a base.
        """
        if self._scratch_base == 0:
            return 0
        # 16-byte alignment
        aligned = (nbytes + 15) & ~15
        addr = self._scratch_base + self._scratch_cursor
        self._scratch_cursor += aligned
        if self._scratch_cursor > self._scratch_size:
            raise RuntimeError(
                f"TLContext scratch overflow: requested {nbytes}B, "
                f"used {self._scratch_cursor}/{self._scratch_size}B"
            )
        return addr

    @property
    def commands(self) -> list[PeCommand]:
        """Return the recorded command trace."""
        return self._commands

    # ── helpers ────────────────────────────────────────────────────

    def _next_handle_id(self) -> str:
        self._handle_counter += 1
        return f"t{self._handle_counter}"

    def _next_completion_id(self) -> str:
        self._completion_counter += 1
        return f"c{self._completion_counter}"

    def _dtype_bytes(self, dtype: str) -> int:
        return _DTYPE_BYTES.get(dtype, 2)

    def _nbytes(self, shape: tuple[int, ...], dtype: str) -> int:
        return math.prod(shape) * self._dtype_bytes(dtype)

    def _charge_dispatch(self, cmd: PeCommand) -> None:
        """Charge structural PE_CPU dispatch cost (ADR-0064 Rev2 D1).

        With a ``cost_model``: emit ``PeCpuOverheadCmd`` carrying
        ``round((FIXED + cmd.logical_bytes × R) / clock_freq_ghz)`` ns just
        before ``cmd``. A CompositeCmd over the descriptor-size cap is
        rejected (D7). Without a cost_model, fall back to the uniform
        ``dispatch_cycles`` contract (ADR-0046 §D6).
        """
        if self._cost_model is not None:
            if isinstance(cmd, CompositeCmd):
                lb = cmd.logical_bytes
                cap = self._cost_model.max_composite_logical_bytes
                if lb > cap:
                    raise ValueError(
                        f"CompositeCmd logical_bytes {lb} exceeds "
                        f"max_composite_logical_bytes {cap} (ADR-0064 D7)"
                    )
            cycles = self._cost_model.dispatch_cycles(
                type(cmd), cmd.logical_bytes)
            ns = round(cycles / self._clock_freq_ghz)
            if ns > 0:
                self._emit(PeCpuOverheadCmd(cycles=ns))
        elif self._dispatch_cycles > 0:
            self._emit(PeCpuOverheadCmd(cycles=self._dispatch_cycles))

    def _make_handle(
        self, addr: int, shape: tuple[int, ...], dtype: str,
        space: str = "tcm", pinned: bool = False,
    ) -> TensorHandle:
        return TensorHandle(
            id=self._next_handle_id(),
            addr=addr, shape=shape, dtype=dtype,
            nbytes=self._nbytes(shape, dtype),
            space=space,
            pinned=pinned,
        )

    def _make_compute_out(
        self, shape: tuple[int, ...], dtype: str,
    ) -> TensorHandle:
        """Allocate an output TensorHandle in PE-local scratch (TCM space).

        Used by math/compute ops so the result has a real address that can
        be the source of a later send/store. The data field stays None in
        Phase 1 — Phase 2 DataExecutor fills the actual ndarray.
        """
        nbytes = self._nbytes(shape, dtype)
        addr = self._scratch_alloc(nbytes)
        return TensorHandle(
            id=self._next_handle_id(),
            addr=addr, shape=shape, dtype=dtype,
            nbytes=nbytes, space="tcm",
            # TCM-resident: a downstream composite operand reads it on-chip,
            # not via a DMA_READ of its (bit-61) scratch address — same as
            # the composite auto-output and recipe scratch handles below.
            pinned=True,
        )

    # ── Reference (no DMA, metadata only) ────────────────────────

    def ref(
        self, ptr: int, shape: tuple[int, ...], dtype: str = "f16",
    ) -> TensorHandle:
        """Create a TensorHandle referencing HBM data without issuing DMA.

        Used when the scheduler will stream data per-tile (e.g., tensor b
        in a composite GEMM) or as a composite's HBM output handle
        (ADR-0065 D8). The handle's ``space="hbm"`` — it references HBM
        data. No command is generated. (Operand-input DMA stays
        ``pinned``-based per ADR-0065 D4, so this label does not change
        input streaming.)
        """
        return self._make_handle(addr=ptr, shape=shape, dtype=dtype,
                                 space="hbm")

    # ── Data Movement (blocking, DMA engine) ──────────────────────

    def _emit(self, cmd: PeCommand) -> Any:
        """Emit command: greenlet switch if runner available, else append to list.

        Each dispatched PeCommand is preceded by a ``PeCpuOverheadCmd``
        carrying its structural dispatch cost (ADR-0064 Rev2 D1).
        ``PeCpuOverheadCmd`` (incl. manual ``tl.cycles``) and ``WaitCmd``
        bypass the charge (D5).
        """
        if not isinstance(cmd, (PeCpuOverheadCmd, WaitCmd)):
            self._charge_dispatch(cmd)
        if self._runner is not None:
            return self._runner.switch_to_simpy(cmd)
        self._commands.append(cmd)
        return None

    def load(
        self, ptr: int, shape: tuple[int, ...], dtype: str = "f16",
    ) -> TensorHandle:
        """Load tensor from HBM — **lazy** (ADR-0062 §D1/§D2).

        Posts the ``DmaReadCmd`` to PE_DMA and returns immediately with a
        ``TensorHandle`` whose ``pending`` field references a fresh
        ``LoadFuture``. The actual DMA completion is awaited at the first
        consuming op (``tl.dot``, MATH, ``tl.store``, ``tl.send``,
        ``tl.copy_to``, ``tl.composite``) via ``_await_pending``.

        Command-list mode: emits the DmaReadCmd to ``self._commands``,
        attaches a LoadFuture for structural compatibility (its event
        stays None — no engine, no SimPy event to wait on).
        """
        nbytes = self._nbytes(shape, dtype)
        # LoadFuture is mutable; create it first, attach to the handle,
        # then point its ``cmd`` at the DmaReadCmd that references the
        # *final* handle. This guarantees ``handle.pending.cmd.handle is handle``.
        future = LoadFuture.__new__(LoadFuture)
        future.event = None
        future.resolved = False
        future.data = None
        handle = TensorHandle(
            id=self._next_handle_id(),
            addr=ptr, shape=shape, dtype=dtype, nbytes=nbytes,
            data=None, space="hbm", pinned=True, pending=future,
        )
        cmd = DmaReadCmd(handle=handle, src_addr=ptr, nbytes=nbytes)
        future.cmd = cmd
        # Lazy load bypasses _emit (it posts via a load_issue future), so it
        # charges its own dispatch cost here (ADR-0064 D1).
        self._charge_dispatch(cmd)
        if self._runner is not None:
            # Lazy: runner posts the DmaReadCmd, sets future.event, then
            # switches back immediately. No yield on completion here.
            self._runner.switch_to_simpy(("load_issue", future))
        else:
            self._commands.append(cmd)
        return handle

    def _await_pending(self, *handles: TensorHandle | None) -> None:
        """Auto-wait at first use (ADR-0062 §D2).

        For each handle carrying an unresolved ``LoadFuture``, yield on
        the DMA completion event (greenlet → runner). Command-list mode
        is a no-op.
        """
        if self._runner is None:
            return
        for h in handles:
            if h is None:
                continue
            pending = getattr(h, "pending", None)
            if pending is None or pending.resolved:
                continue
            if isinstance(pending, CompositeFuture):
                # Composite output: wait on its completion (ADR-0065 D8).
                self._emit(WaitCmd(handle=pending.completion))
                pending.resolved = True
            else:
                self._runner.switch_to_simpy(("load_await", h))

    def store(self, ptr: int, handle: TensorHandle) -> None:
        """Store tensor from TCM to HBM."""
        self._await_pending(handle)
        cmd = DmaWriteCmd(handle=handle, dst_addr=ptr, nbytes=handle.nbytes)
        self._emit(cmd)

    def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None:
        """TCM-to-TCM byte copy (ADR-0063 §D3.1).

        Writes ``src``'s bytes into ``dst``'s address. Both handles must
        live in TCM and have matching shape and dtype. Symmetric to
        ``tl.store`` (which targets HBM) but stays on-chip so it doesn't
        emit a DMA entry into op_log.

        Used inside ``tl.scratch_scope()`` to persist a scoped result —
        typically an updated running ``(m, ℓ, O)`` — to an outside-scope
        (persistent) handle so its bytes survive the scope's ``__exit__``
        cursor rewind (ADR-0063 §D3 two-arena pattern).
        """
        if src.shape != dst.shape:
            raise ValueError(
                f"tl.copy_to: shape mismatch — src.shape={src.shape} "
                f"vs dst.shape={dst.shape}"
            )
        if src.dtype != dst.dtype:
            raise ValueError(
                f"tl.copy_to: dtype mismatch — src.dtype={src.dtype!r} "
                f"vs dst.dtype={dst.dtype!r}"
            )
        if dst.space != "tcm":
            raise ValueError(
                f"tl.copy_to: dst must be in TCM (got space={dst.space!r}); "
                "writes to HBM go through tl.store"
            )
        if src.space != "tcm":
            raise ValueError(
                f"tl.copy_to: src must be in TCM (got space={src.space!r}); "
                "reads from HBM go through tl.load"
            )
        self._await_pending(src)
        self._emit(CopyCmd(src=src, dst=dst, nbytes=src.nbytes))

    # ── GEMM Engine (blocking) ────────────────────────────────────

    def dot(self, a: TensorHandle, b: TensorHandle) -> TensorHandle:
        """Matrix multiply: out = a @ b. Both operands must be in TCM.

        a: (M, K), b: (K, N) → out: (M, N)
        """
        if len(a.shape) < 2 or len(b.shape) < 2:
            raise ValueError("dot requires 2D tensors")
        m, k = a.shape[-2], a.shape[-1]
        k2, n = b.shape[-2], b.shape[-1]
        if k != k2:
            raise ValueError(f"dot shape mismatch: a.K={k} != b.K={k2}")
        out_shape = (*a.shape[:-2], m, n)
        out_dtype = a.dtype
        out = self._make_compute_out(shape=out_shape, dtype=out_dtype)
        self._await_pending(a, b)
        self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n))
        return out

    # ── MATH Engine: unary (blocking) ─────────────────────────────

    def _unary_math(self, op: str, x: TensorHandle) -> TensorHandle:
        out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
        self._await_pending(x)
        self._emit(MathCmd(op=op, inputs=(x,), out=out))
        return out

    def exp(self, x: TensorHandle) -> TensorHandle:
        return self._unary_math("exp", x)

    def log(self, x: TensorHandle) -> TensorHandle:
        return self._unary_math("log", x)

    def sqrt(self, x: TensorHandle) -> TensorHandle:
        return self._unary_math("sqrt", x)

    def abs(self, x: TensorHandle) -> TensorHandle:
        return self._unary_math("abs", x)

    def sigmoid(self, x: TensorHandle) -> TensorHandle:
        return self._unary_math("sigmoid", x)

    def cos(self, x: TensorHandle) -> TensorHandle:
        return self._unary_math("cos", x)

    def sin(self, x: TensorHandle) -> TensorHandle:
        return self._unary_math("sin", x)

    # ── MATH Engine: reduction (blocking) ─────────────────────────

    def _reduction(
        self, op: str, x: TensorHandle, axis: int,
    ) -> TensorHandle:
        out_shape = list(x.shape)
        out_shape[axis] = 1
        out = self._make_compute_out(shape=tuple(out_shape), dtype=x.dtype)
        self._await_pending(x)
        self._emit(MathCmd(op=op, inputs=(x,), out=out, axis=axis))
        return out

    def sum(self, x: TensorHandle, axis: int) -> TensorHandle:
        return self._reduction("sum", x, axis)

    def max(self, x: TensorHandle, axis: int) -> TensorHandle:
        return self._reduction("max", x, axis)

    def min(self, x: TensorHandle, axis: int) -> TensorHandle:
        return self._reduction("min", x, axis)

    # ── MATH Engine: binary (blocking) ────────────────────────────

    def _binary_math(
        self, op: str, a: TensorHandle, b: TensorHandle,
    ) -> TensorHandle:
        out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
        self._await_pending(a, b)
        self._emit(MathCmd(op=op, inputs=(a, b), out=out))
        return out

    def where(
        self, cond: TensorHandle, a: TensorHandle, b: TensorHandle,
    ) -> TensorHandle:
        out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
        self._await_pending(cond, a, b)
        self._emit(MathCmd(op="where", inputs=(cond, a, b), out=out))
        return out

    def maximum(self, a: TensorHandle, b: TensorHandle) -> TensorHandle:
        """Element-wise max of two tensors (real Triton: tl.maximum)."""
        return self._binary_math("maximum", a, b)

    def minimum(self, a: TensorHandle, b: TensorHandle) -> TensorHandle:
        """Element-wise min of two tensors (real Triton: tl.minimum)."""
        return self._binary_math("minimum", a, b)

    def fma(
        self, a: TensorHandle, b: TensorHandle, c: TensorHandle,
    ) -> TensorHandle:
        """Fused multiply-add: a * b + c (real Triton: tl.fma)."""
        out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
        self._await_pending(a, b, c)
        self._emit(MathCmd(op="fma", inputs=(a, b, c), out=out))
        return out

    def clamp(
        self,
        x: TensorHandle,
        min: TensorHandle,
        max: TensorHandle,
    ) -> TensorHandle:
        """Clamp x to [min, max] (real Triton: tl.clamp)."""
        out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
        self._await_pending(x, min, max)
        self._emit(MathCmd(op="clamp", inputs=(x, min, max), out=out))
        return out

    def softmax(self, x: TensorHandle, axis: int = -1) -> TensorHandle:
        """Numerically-stable softmax along ``axis`` (real Triton: tl.softmax).

        Implemented as a single MathCmd (op="softmax") so timing accounts
        for one MATH dispatch; Phase 2 DataExecutor expands it to the
        canonical (x - max) → exp → sum → div sequence.
        """
        out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
        self._await_pending(x)
        self._emit(MathCmd(op="softmax", inputs=(x,), out=out, axis=axis))
        return out

    # ── Scalar helpers (real Triton: tl.cdiv etc.) ────────────────

    @staticmethod
    def cdiv(a: int, b: int) -> int:
        """Ceiling division: (a + b - 1) // b (real Triton: tl.cdiv).

        Used by host/kernel grid math; not a tensor op, so no MathCmd
        is emitted. Mirrors triton.cdiv.
        """
        return -(-int(a) // int(b))

    # ── Index / Scalar (PE_CPU, no engine) ────────────────────────

    def program_id(self, axis: int = 0) -> int:
        """Return program instance index (ADR-0022).

        axis=0: local PE id within cube.
        axis=1: cube id.
        """
        if axis == 1:
            return self._cube_id
        return self._pe_id

    def num_programs(self, axis: int = 0) -> int:
        """Return total number of program instances (ADR-0022).

        axis=0: num PEs per cube.
        axis=1: num cubes.
        """
        if axis == 1:
            return self._num_cubes
        return self._num_programs

    def arange(self, start: int, end: int, dtype: str = "i32") -> TensorHandle:
        """Create index range tensor in TCM."""
        n = end - start
        return self._make_handle(addr=0, shape=(n,), dtype=dtype)

    def zeros(self, shape: tuple[int, ...], dtype: str = "f16") -> TensorHandle:
        """Create zero-filled tensor in TCM."""
        return self._make_handle(addr=0, shape=shape, dtype=dtype)

    def full(
        self, shape: tuple[int, ...], value: float | int, dtype: str = "f16",
    ) -> TensorHandle:
        """Create constant-filled tensor in TCM."""
        return self._make_handle(addr=0, shape=shape, dtype=dtype)

    # ── Metadata (no compute, no DMA) ─────────────────────────────

    def trans(self, x: TensorHandle) -> TensorHandle:
        """Transpose — shape change only, no command generated."""
        if len(x.shape) < 2:
            raise ValueError("trans requires at least 2D tensor")
        new_shape = (*x.shape[:-2], x.shape[-1], x.shape[-2])
        return TensorHandle(
            id=x.id, addr=x.addr, shape=new_shape,
            dtype=x.dtype, nbytes=x.nbytes, data=x.data,
        )

    # ── IPCQ (CCL) collective primitives (ADR-0023 D4) ────────────

    def send(
        self,
        dir: str,
        src: TensorHandle | None = None,
        *,
        src_addr: int | None = None,
        nbytes: int | None = None,
        shape: tuple[int, ...] | None = None,
        dtype: str = "f16",
        space: str = "tcm",
    ) -> None:
        """Send tensor data to the peer in the given direction.

        Two calling forms:
            tl.send(dir, handle)                       # use handle's metadata
            tl.send(dir, src_addr=..., nbytes=..., shape=..., dtype=..., space=...)

        Blocking: returns when PE_IPCQ has accepted the request and
        forwarded the IpcqDmaToken to PE_DMA. Backpressure may apply.
        """
        if src is not None:
            src_addr = src.addr
            nbytes = src.nbytes
            shape = src.shape
            dtype = src.dtype
            space = getattr(src, "space", space)
        if src_addr is None or nbytes is None or shape is None:
            raise ValueError("tl.send: provide either a TensorHandle or src_addr/nbytes/shape")
        # ADR-0062: if the source is a lazy-loaded handle, await first so
        # the data snapshot below sees the real bytes.
        self._await_pending(src)
        # Carry the handle's .data snapshot (if available). When the source
        # is a recv slot, .data holds the numpy array that was read from
        # MemoryStore at recv-time. This prevents a Phase 1 race where a
        # later IPCQ inbound overwrites the slot before the outbound
        # PE_DMA reads it.
        handle_data = getattr(src, "data", None) if src is not None else None
        cmd = IpcqSendCmd(
            direction=dir,
            src_addr=src_addr, src_space=space,
            nbytes=nbytes, shape=shape, dtype=dtype,
            handle_id=self._next_handle_id(),
            data=handle_data,
        )
        self._emit(cmd)

    def recv(
        self,
        dir: str | None = None,
        shape: tuple[int, ...] = (),
        dtype: str = "f16",
        space: str = "tcm",
        dst_addr: int | None = None,
        dst_space: str | None = None,
    ) -> TensorHandle:
        """Receive tensor data from a peer.

        Args:
            dir: specific direction (e.g. "W"), or None for round-robin.
            shape, dtype: expected tensor metadata.
            dst_addr / dst_space: if both are provided, the slot data is
                copied to (dst_space, dst_addr) before the handle is
                returned ("copy_to_dst" mode). Otherwise the slot address
                is returned directly ("return_slot" mode).

        Returns:
            TensorHandle pointing to the slot (or dst) where the data has
            arrived. In greenlet/runner mode, ``handle.data`` carries the
            actual ndarray; in command-list mode the handle is a placeholder.
        """
        if dst_addr is not None and dst_space is not None:
            cmd = IpcqRecvCmd(
                direction=dir,
                shape=shape, dtype=dtype,
                handle_id=self._next_handle_id(),
                recv_mode="copy_to_dst",
                dst_addr=dst_addr, dst_space=dst_space,
            )
        else:
            cmd = IpcqRecvCmd(
                direction=dir,
                shape=shape, dtype=dtype,
                handle_id=self._next_handle_id(),
            )
        result = self._emit(cmd)
        if isinstance(result, dict):
            slot_addr = int(result.get("src_addr", 0))
            slot_space = str(result.get("src_space", "tcm"))
            data = result.get("data")
            return TensorHandle(
                id=self._next_handle_id(),
                addr=slot_addr,
                shape=shape,
                dtype=dtype,
                nbytes=self._nbytes(shape, dtype),
                data=data,
                space=slot_space,
                # On-chip (tcm/sram) IPCQ slot: read in place as a composite
                # operand, not DMA_READ of its non-HBM slot address. An HBM
                # slot is a valid PA, so it stays unpinned (DMA streams it).
                pinned=slot_space != "hbm",
            )
        return self._make_handle(addr=0, shape=shape, dtype=dtype)

    def recv_no_consume(
        self,
        dir: str | None = None,
        shape: tuple[int, ...] = (),
        dtype: str = "f16",
    ) -> TensorHandle:
        """DIAGNOSTIC ONLY — recv that blocks for arrival but skips slot read.

        Same blocking semantics as ``tl.recv``: the kernel waits until
        the payload has landed in the IPCQ slot. Differs from ``tl.recv``
        by skipping the slot-read latency charge (slot-IO + PE↔bank
        fabric drain) on DST.

        This entry point exists solely so the pe2pe overview plot can
        draw an apples-to-apples comparison against ``tl.store`` (a
        one-sided fabric write that pays no read on DST). Production
        kernels MUST use ``tl.recv`` — they need to consume the data
        they receive. This API is segregated from ``tl.recv`` so the
        diagnostic flag can never accidentally be set in real workloads.
        """
        cmd = IpcqRecvCmd(
            direction=dir,
            shape=shape, dtype=dtype,
            handle_id=self._next_handle_id(),
            consume=False,
        )
        result = self._emit(cmd)  # type: ignore[arg-type]
        if isinstance(result, dict):
            slot_addr = int(result.get("src_addr", 0))
            slot_space = str(result.get("src_space", "tcm"))
            return TensorHandle(
                id=self._next_handle_id(),
                addr=slot_addr,
                shape=shape,
                dtype=dtype,
                nbytes=self._nbytes(shape, dtype),
                data=None,
                space=slot_space,
                pinned=slot_space != "hbm",
            )
        return self._make_handle(addr=0, shape=shape, dtype=dtype)

    def recv_async(
        self,
        dir: str,
        shape: tuple[int, ...] = (),
        dtype: str = "f16",
    ) -> "RecvFuture":
        """Non-blocking recv. Returns a future to pass into ``tl.wait``."""
        cmd = IpcqRecvCmd(
            direction=dir,
            shape=shape, dtype=dtype,
            handle_id=self._next_handle_id(),
            blocking=False,
        )
        # recv_async bypasses _emit (posts via a recv_async future), so it
        # charges its own dispatch cost here (ADR-0064 D1).
        self._charge_dispatch(cmd)
        future = RecvFuture(cmd=cmd)
        if self._runner is not None:
            self._runner.switch_to_simpy(("recv_async", future))
        return future

    # ── Composite + Control ───────────────────────────────────────

    def composite(
        self,
        op: Literal["gemm", "math"],
        a: TensorHandle | None = None,
        b: TensorHandle | None = None,
        math_op: str | None = None,
        *,
        out: TensorHandle | None = None,
        out_ptr: int | None = None,
        prologue: list[dict] | None = None,
        epilogue: list[dict] | None = None,
        acc_dtype: str | None = None,
        tile_shape: tuple[int, int, int] | None = None,
    ) -> TensorHandle:
        """Submit a composite command (non-blocking, tiled pipeline).

        ``prologue`` (ADR-0065 D5) is an ordered list of recipe dicts — each
        with an ``"op"`` naming a ``RECIPE_DESCRIPTORS`` entry plus its
        operands. TLContext expands each into flat MATH OpSpecs that run
        before the head op and auto-binds the recipe's primary output into
        the head GEMM's ``a`` operand (D6.6). ``epilogue`` is an ordered list
        of ``EPILOGUE_OPS`` dicts.

        ``out`` (ADR-0065 D8) is a ``TensorHandle``: an HBM ref
        (``tl.ref(addr, shape)`` → ``space="hbm"`` → DMA_WRITE) or an
        in-place TCM handle (e.g. a recipe accumulator → STORE only).
        ``out_ptr`` (int) is shorthand for an HBM output — equivalent to
        ``out=tl.ref(out_ptr, <output shape>)``. If neither is given, a TCM
        scratch is auto-allocated (chainable, like ``tl.dot``).

        Returns the **output TensorHandle** — auto-awaited by downstream ops
        (and by ``tl.wait(handle)``) via its ``pending`` field.
        """
        # ADR-0062: composite operand DMA paths still need their inputs
        # to be resolved before the composite reads them via PE_SCHEDULER.
        self._await_pending(a, b)

        # Prologue recipes → flat MATH OpSpecs + primary-output handle (D5).
        prologue_ops: tuple[OpSpec, ...] = ()
        primary_out: TensorHandle | None = None
        rw_handles: tuple[TensorHandle, ...] = ()
        if prologue:
            prologue_ops, primary_out, rw_handles = self._expand_prologue(prologue)
            # Await any prologue operand that is a pending composite output
            # (e.g. the score tile produced by a prior #1 composite) so it is
            # computed before this composite reads it (RAW across composites).
            for _item in prologue:
                self._await_pending(
                    *[v for v in _item.values() if isinstance(v, TensorHandle)]
                )

        # Auto-bind (D6.6): the recipe's primary output fills the head GEMM's
        # `a` unless the kernel supplied `a` explicitly (then it's ambiguous).
        if op == "gemm" and primary_out is not None:
            if a is not None:
                raise ValueError(
                    "auto-bind conflict: head GEMM operand 'a' was provided "
                    "explicitly but the prologue recipe declares a primary_out "
                    "(ADR-0065 D6.6) — omit 'a' to let the recipe bind it"
                )
            a = primary_out

        if a is None:
            raise ValueError(
                "composite requires operand 'a' (or a prologue recipe with a "
                "primary_out)"
            )

        # Output geometry.
        if op == "gemm" and b is not None:
            m, k = a.shape[-2], a.shape[-1]
            n = b.shape[-1]
            out_dtype = a.dtype
            out_shape: tuple[int, ...] = (m, n)
            out_nbytes = m * n * self._dtype_bytes(out_dtype)
        else:
            out_dtype = a.dtype
            out_shape = a.shape
            out_nbytes = a.nbytes

        # Write-back handle (ADR-0065 D8). Explicit `out` (HBM ref or an
        # in-place TCM accumulator) wins; `out_ptr` is HBM shorthand; else
        # auto-allocate a TCM scratch — the chainable default, like tl.dot.
        # pinned=True on the auto output: it is TCM-resident, so a downstream
        # op reading it needs no DMA_READ.
        if out is not None:
            out_handle = out
        elif out_ptr is not None:
            out_handle = TensorHandle(
                id=self._next_handle_id(),
                addr=out_ptr, shape=out_shape, dtype=out_dtype,
                nbytes=out_nbytes, space="hbm",
            )
        else:
            out_handle = TensorHandle(
                id=self._next_handle_id(),
                addr=self._scratch_alloc(out_nbytes),
                shape=out_shape, dtype=out_dtype, nbytes=out_nbytes,
                space="tcm", pinned=True,
            )

        # Head op (flat-ops, ADR-0065 D2): GEMM takes named a/b and carries
        # m/k/n in extra; MATH takes named a and carries math_op in extra.
        if op == "gemm" and b is not None:
            head_operands: dict[str, Any] = {"a": a, "b": b}
            head_extra: dict[str, Any] = {"m": m, "k": k, "n": n}
        else:
            head_operands = {"a": a}
            head_extra = {}
        for _k, _v in (("acc_dtype", acc_dtype),
                       ("tile_shape", tile_shape),
                       ("math_op", math_op)):
            if _v is not None:
                head_extra[_k] = _v

        head_spec = OpSpec(
            kind=op, scope=Scope.OUTPUT_TILE,
            operands=head_operands, extra=head_extra, out=out_handle,
        )
        epi_specs = tuple(self._build_epilogue_spec(e, i)
                          for i, e in enumerate(epilogue or []))
        ops_tuple = (*prologue_ops, head_spec, *epi_specs)

        self._validate_composite_ops(ops_tuple, prologue_ops)

        completion = CompletionHandle(id=self._next_completion_id())
        self._emit(CompositeCmd(
            completion=completion, ops=ops_tuple, rw_handles=rw_handles,
        ))
        # Attach the completion to the output handle so downstream ops (and
        # tl.wait) auto-await it (ADR-0065 D8 / ADR-0062 lazy pattern), and
        # return it so the result chains like tl.dot's (frozen dataclass →
        # set pending via object.__setattr__).
        object.__setattr__(out_handle, "pending", CompositeFuture(completion))
        return out_handle

    def _expand_prologue(
        self, prologue: list[dict],
    ) -> tuple[tuple[OpSpec, ...], TensorHandle | None, tuple[TensorHandle, ...]]:
        """Expand prologue recipe items into flat MATH OpSpecs (ADR-0065 D5).

        Returns ``(math_ops, primary_out_handle, rw_handles)``. PE_SCHEDULER
        never sees the recipe — only the flat ops.
        """
        from kernbench.triton_emu.tl_recipes import (
            PRIMARY_OUT_SLOT, RECIPE_DESCRIPTORS,
        )

        math_ops: list[OpSpec] = []
        primary_out: TensorHandle | None = None
        rw_handles: list[TensorHandle] = []

        for item in prologue:
            if not isinstance(item, dict) or "op" not in item:
                raise ValueError("prologue entry must be a dict with an 'op' key")
            name = item["op"]
            recipe = RECIPE_DESCRIPTORS.get(name)
            if recipe is None:
                known = ", ".join(sorted(RECIPE_DESCRIPTORS))
                raise ValueError(
                    f"unknown prologue recipe {name!r} (known: {known})"
                )

            # Bind recipe operands from the item.
            slots: dict[str, TensorHandle] = {}
            for opd_name in recipe.operands:
                if opd_name not in item:
                    raise ValueError(
                        f"prologue recipe {name!r} missing operand {opd_name!r}"
                    )
                slots[opd_name] = item[opd_name]

            # Expand engine_seq sequentially — later ops read earlier dsts.
            for eop in recipe.engine_seq:
                operands: dict[str, Any] = {}
                for key in ("src", "src_a", "src_b", "src_c"):
                    slot_name = getattr(eop, key)
                    if slot_name is None:
                        continue
                    if slot_name not in slots:
                        raise ValueError(
                            f"recipe {name!r} op {eop.op_kind!r} reads unbound "
                            f"slot {slot_name!r}"
                        )
                    operands[key] = slots[slot_name]

                dst_handle = self._resolve_recipe_dst(eop, recipe, slots, operands)
                extra: dict[str, Any] = {}
                if eop.reduce_axis is not None:
                    extra["reduce_axis"] = eop.reduce_axis
                if eop.bcast_axis is not None:
                    extra["bcast_axis"] = eop.bcast_axis

                math_ops.append(OpSpec(
                    kind=eop.op_kind, scope=Scope.KERNEL,
                    operands=operands, extra=extra, out=dst_handle,
                ))
                slots[eop.dst] = dst_handle

            if recipe.primary_out is not None:
                primary_out = slots.get(PRIMARY_OUT_SLOT)
            for opd_name, rw in recipe.operands.items():
                if rw == "RW":
                    rw_handles.append(slots[opd_name])

        return tuple(math_ops), primary_out, tuple(rw_handles)

    def _resolve_recipe_dst(
        self, eop: Any, recipe: Any, slots: dict, operands: dict,
    ) -> TensorHandle:
        """Resolve an EngineOp's dst to a handle: reuse an existing slot
        (RW operand / already-allocated), else allocate fresh TCM scratch
        with an inferred shape."""
        from kernbench.triton_emu.tl_recipes import PRIMARY_OUT_SLOT

        dst_name = eop.dst
        if dst_name in slots:
            return slots[dst_name]            # write in place

        if recipe.primary_out is not None and dst_name == PRIMARY_OUT_SLOT:
            ref = slots[recipe.primary_out.from_shape]
            shape, dtype = ref.shape, ref.dtype
        elif eop.reduce_axis is not None:
            # Reductions keep the reduced axis as size 1 (keepdims), so the
            # result broadcasts against the running (m, l) — e.g. rmax of
            # s=(G,TILE) over axis -1 → (G,1), not (G,) (ADR-0065 N2).
            ref = next(iter(operands.values()))
            dims = list(ref.shape)
            if len(dims) > 1:
                ax = eop.reduce_axis % len(dims)
                dims[ax] = 1
            shape = tuple(dims)
            dtype = ref.dtype
        else:
            ref = next(iter(operands.values()))
            shape, dtype = ref.shape, ref.dtype

        nbytes = self._nbytes(shape, dtype)
        addr = self._scratch_alloc(nbytes)
        # pinned=True: recipe scratch / primary-out live in TCM already, so
        # the head GEMM's auto-bound `a` (= primary-out P) is consumed in
        # place — no DMA_READ (ADR-0065 P3, pinned-based DMA decision).
        return TensorHandle(
            id=self._next_handle_id(),
            addr=addr, shape=shape, dtype=dtype, nbytes=nbytes,
            space="tcm", pinned=True,
        )

    @staticmethod
    def _validate_composite_ops(
        ops: tuple[OpSpec, ...], prologue_ops: tuple[OpSpec, ...],
    ) -> None:
        """Emit-time invariants (ADR-0065 D6.1 / D6.7).

        D6.1 (GEMM count ≤ 1) applies to the whole composite. D6.7 (MATH
        operands TCM-only) applies to **prologue recipe ops only** — the
        head op (gemm *or* math) and epilogue ops keep the existing
        DMA-staged-from-HBM behavior (a math head's `a` is streamed by
        PE_SCHEDULER, like a GEMM operand).
        """
        gemm_count = sum(1 for o in ops if o.kind == "gemm")
        if gemm_count > 1:
            raise ValueError(
                f"composite must carry at most 1 GEMM op (ADR-0065 D6.1); "
                f"got {gemm_count}"
            )
        for o in prologue_ops:
            for h in o.operands.values():
                if getattr(h, "space", "tcm") != "tcm":
                    raise ValueError(
                        f"MATH operand must be TCM-resident (ADR-0065 D6.7); "
                        f"prologue op {o.kind!r} operand has space="
                        f"{getattr(h, 'space', None)!r}"
                    )

    @staticmethod
    def _build_epilogue_spec(entry: dict, idx: int) -> OpSpec:
        if not isinstance(entry, dict) or "op" not in entry:
            raise ValueError(
                f"epilogue[{idx}]: each entry must be a dict with an 'op' key"
            )
        kind = entry["op"]
        if kind not in EPILOGUE_OPS:
            known = ", ".join(sorted(EPILOGUE_OPS))
            raise ValueError(
                f"epilogue[{idx}]: unknown op {kind!r} "
                f"(known ops: {known})"
            )
        required, default_scope = EPILOGUE_OPS[kind]
        missing = [f for f in required if f not in entry]
        if missing:
            raise ValueError(
                f"epilogue[{idx}] op {kind!r} missing required field(s): "
                f"{', '.join(missing)}"
            )
        scope = Scope(entry["scope"]) if "scope" in entry else default_scope
        operands: dict = {}
        scalar: float | None = None
        extra: dict = {}
        for f in required:
            v = entry[f]
            if isinstance(v, TensorHandle):
                operands[f] = v
            elif isinstance(v, (int, float)):
                if scalar is None:
                    scalar = float(v)
                else:
                    extra[f] = v
            else:
                extra[f] = v
        return OpSpec(
            kind=kind, scope=scope,
            operands=operands, scalar=scalar, extra=extra,
        )

    def wait(self, handle: "CompletionHandle | RecvFuture | None" = None) -> Any:
        """Wait for a composite, a recv future, or all pending composites.

        - ``CompletionHandle`` (or None): wait for composite completion.
        - ``RecvFuture``: wait for a non-blocking ``recv_async`` to finish.
          Returns the resolved ``TensorHandle``.
        """
        if isinstance(handle, RecvFuture):
            if handle.resolved:
                return handle.result
            if self._runner is None:
                raise RuntimeError(
                    "tl.wait(RecvFuture) requires runner mode (greenlet)"
                )
            result_dict = self._runner.switch_to_simpy(("recv_wait", handle))
            slot_addr = int(result_dict.get("src_addr", 0))
            slot_space = str(result_dict.get("src_space", "tcm"))
            data = result_dict.get("data")
            th = TensorHandle(
                id=self._next_handle_id(),
                addr=slot_addr,
                shape=handle.cmd.shape,
                dtype=handle.cmd.dtype,
                nbytes=self._nbytes(handle.cmd.shape, handle.cmd.dtype),
                data=data,
                space=slot_space,
                pinned=slot_space != "hbm",
            )
            handle.resolved = True
            handle.result = th
            return th

        # Composite output handle (ADR-0065 D8): wait on its completion.
        if isinstance(handle, TensorHandle):
            pending = getattr(handle, "pending", None)
            if isinstance(pending, CompositeFuture):
                if not pending.resolved:
                    self._emit(WaitCmd(handle=pending.completion))
                    pending.resolved = True
            return None

        # Composite path (CompletionHandle or None — existing behaviour).
        self._emit(WaitCmd(handle=handle))
        return None

    def cycles(self, n: int) -> None:
        """Declare PE_CPU scalar execution overhead (cycles)."""
        self._emit(PeCpuOverheadCmd(cycles=n))


# ── TensorHandle arithmetic operators ─────────────────────────────
# Enables: a + b, a * b, a - b, a / b in kernel code.
# Each creates a MathCmd via a module-level helper that requires a
# TLContext. We attach the context to handles via a closure approach.


def _enable_tensor_ops() -> None:
    """Patch TensorHandle with arithmetic operators.

    Called once at module load. Operators create MathCmd entries via
    a thread-local TLContext reference set during kernel execution.
    """
    import threading

    _local = threading.local()

    def set_active_context(ctx: TLContext | None) -> None:
        _local.ctx = ctx

    def get_active_context() -> TLContext:
        ctx = getattr(_local, "ctx", None)
        if ctx is None:
            raise RuntimeError("TensorHandle ops require an active TLContext")
        return ctx

    def _binop(op: str):
        def method(self: TensorHandle, other: TensorHandle) -> TensorHandle:
            ctx = get_active_context()
            return ctx._binary_math(op, self, other)
        return method

    # Patch TensorHandle class with operators
    TensorHandle.__add__ = _binop("add")       # type: ignore[attr-defined]
    TensorHandle.__sub__ = _binop("sub")       # type: ignore[attr-defined]
    TensorHandle.__mul__ = _binop("mul")       # type: ignore[attr-defined]
    TensorHandle.__truediv__ = _binop("div")   # type: ignore[attr-defined]

    # Expose context management
    TLContext._set_active = staticmethod(set_active_context)  # type: ignore[attr-defined]
    TLContext._get_active = staticmethod(get_active_context)  # type: ignore[attr-defined]


_enable_tensor_ops()


def run_kernel(
    kernel_fn,
    tl_ctx: TLContext,
    *args,
    **kwargs,
) -> list[PeCommand]:
    """Execute a kernel function with the given TLContext and return commands.

    Sets tl_ctx as the active context for TensorHandle operators,
    calls the kernel, then clears the context.
    """
    TLContext._set_active(tl_ctx)  # type: ignore[attr-defined]
    try:
        kernel_fn(*args, tl=tl_ctx, **kwargs)
    finally:
        TLContext._set_active(None)  # type: ignore[attr-defined]
    return tl_ctx.commands
