"""Regression: a recv handle used as a composite ``b`` operand must be read
in place (on-chip), NOT DMA_READ from its non-HBM slot address.

Before the recv-pinning fix, recv handles were unpinned → the composite
scheduler emitted DMA_READ(B) from the IPCQ slot address (non-HBM) →
PhysAddr.decode raised PhysAddrError. On-chip (TCM/SRAM) recv slots are
now pinned (ADR-0065 D4), so a composite operand reads them in place.

Isolates operand ``b``: ``a`` is a pinned ``tl.load``; only ``b`` is the
recv handle. Mirrors the decode row-chain IPCQ directions (intra_E recv /
intra_W send) so the recv actually lands without deadlock.
"""
from __future__ import annotations

from pathlib import Path

from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
from kernbench.policy.placement.dp import DPPolicy
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology

TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[1] / "topology.yaml"
P = 8


def _ccl_cfg():
    return resolve_algorithm_config(
        load_ccl_config(), name="lrab_hierarchical_allreduce",
    )


def _engine_factory(t, d):
    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)


def _recv_operand_kernel(a_ptr, b_ptr, o_ptr, M, K, N, P, *, tl):
    pe_id = tl.program_id(axis=0)
    pe_in_group = pe_id % P
    group_cols = min(4, P)
    pe_col = pe_in_group % group_cols

    A = tl.load(a_ptr, shape=(M, K), dtype="f16")          # pinned (no DMA)
    if pe_col < group_cols - 1:
        B = tl.recv(dir="intra_E", shape=(K, N), dtype="f16")  # recv handle
        out = tl.composite("gemm", a=A, b=B)                   # ← b = recv slot
        tl.store(o_ptr, out)
    if pe_col > 0:
        B_send = tl.load(b_ptr, shape=(K, N), dtype="f16")
        tl.send(dir="intra_W", src=B_send)


def test_recv_handle_as_composite_operand_reads_in_place():
    M = K = N = 64
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
        dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=P)
        a = ctx.zeros((M, K), dtype="f16", dp=dp, name="recv_operand_a")
        b = ctx.zeros((K, N), dtype="f16", dp=dp, name="recv_operand_b")
        o = ctx.empty((M, N), dtype="f16", dp=dp, name="recv_operand_o")
        ctx.launch(
            "recv_operand_composite", _recv_operand_kernel,
            a, b, o, M, K, N, P,
            _auto_dim_remap=False,
        )

    r = run_bench(
        topology=topo, bench_fn=_bench_fn,
        device=resolve_device(None), engine_factory=_engine_factory,
    )
    assert r.completion.ok, f"recv-operand composite completion: {r.completion}"