"""Phase 1 spec tests for ADR-0065 D8 — composite output handle + output-space DMA.

`tl.composite(...)` returns the **output TensorHandle** (not a bare
CompletionHandle) so its result chains into downstream ops like `tl.dot`'s.
`out` is always a handle: `out=tl.ref(addr, shape)` for HBM output (DMA_WRITE
inside the composite), `out=<TCM handle>` for in-place, or omitted → TLContext
auto-allocates a TCM scratch. The output handle's `space` decides the tile
loop's write-back (tcm → STORE only; hbm → STORE + DMA_WRITE). `tl.ref`
returns `space="hbm"`.

Phase 1 (this commit): tests only. FAIL until D8:
  - `tl.ref` returns space="tcm"; `composite` returns CompletionHandle;
    output always DMA_WRITEs; `tl.wait` rejects a TensorHandle.
"""
from __future__ import annotations

from kernbench.common.pe_commands import OpSpec, Scope, TensorHandle
from kernbench.triton_emu.tl_context import TLContext


def _tl() -> TLContext:
    return TLContext(pe_id=0, num_programs=1, scratch_base=0x200000,
                     scratch_size=1 << 20)


def test_ref_returns_hbm_space():
    h = _tl().ref(0x1000, shape=(4, 4), dtype="f16")
    assert h.space == "hbm"


def test_composite_returns_handle_auto_tcm():
    tl = _tl()
    a = tl.ref(0x1000, shape=(8, 16), dtype="f16")
    b = tl.ref(0x2000, shape=(16, 8), dtype="f16")
    out = tl.composite(op="gemm", a=a, b=b)        # out omitted → auto TCM
    assert isinstance(out, TensorHandle)
    assert out.space == "tcm"
    assert out.addr >= 0x200000                    # auto-allocated scratch
    assert out.shape == (8, 8)


def test_composite_returns_the_explicit_out_handle():
    tl = _tl()
    a = tl.ref(0x1000, shape=(8, 16), dtype="f16")
    b = tl.ref(0x2000, shape=(16, 8), dtype="f16")
    O = TensorHandle(id="O", addr=0x5000, shape=(8, 8), dtype="f16",
                     nbytes=128, space="tcm", pinned=True)
    out = tl.composite(op="gemm", a=a, b=b, out=O)
    assert out is O


def test_composite_output_chains_into_next_op():
    tl = _tl()
    a = tl.ref(0x1000, shape=(8, 16), dtype="f16")
    b = tl.ref(0x2000, shape=(16, 8), dtype="f16")
    scores = tl.composite(op="gemm", a=a, b=b)     # handle
    out2 = tl.composite(op="math", a=scores, math_op="exp")
    assert isinstance(out2, TensorHandle)


def test_wait_accepts_tensor_handle():
    tl = _tl()
    a = tl.ref(0x1000, shape=(8, 16), dtype="f16")
    b = tl.ref(0x2000, shape=(16, 8), dtype="f16")
    out = tl.composite(op="gemm", a=a, b=b)
    tl.wait(out)                                   # must not raise


def test_output_space_drives_dma_write():
    from kernbench.components.builtin.pe_types import StageType
    from kernbench.components.builtin.tiling import generate_plan_from_ops

    a = TensorHandle(id="a", addr=0x1000, shape=(64, 64), dtype="f16",
                     nbytes=8192, space="tcm", pinned=True)
    b = TensorHandle(id="b", addr=0x2000, shape=(64, 64), dtype="f16",
                     nbytes=8192, space="hbm")

    def _gemm(out):
        return OpSpec(kind="gemm", scope=Scope.OUTPUT_TILE,
                      operands={"a": a, "b": b}, out=out,
                      extra={"m": 64, "k": 64, "n": 64})

    o_tcm = TensorHandle(id="ot", addr=0x300000, shape=(64, 64), dtype="f16",
                         nbytes=8192, space="tcm", pinned=True)
    o_hbm = TensorHandle(id="oh", addr=0x9000, shape=(64, 64), dtype="f16",
                         nbytes=8192, space="hbm")

    def _has_write(ops):
        plan = generate_plan_from_ops(ops, tile_m=64, tile_k=64, tile_n=64,
                                      bytes_per_element=2,
                                      pe_prefix="sip0.cube0.pe0")
        return any(s.stage_type == StageType.DMA_WRITE
                   for s in plan.tiles[0].stages)

    assert not _has_write((_gemm(o_tcm),)), "tcm out → no DMA_WRITE"
    assert _has_write((_gemm(o_hbm),)), "hbm out → DMA_WRITE"