"""Test that tl.recv() (no direction) works under the mock runtime
and the SimPy PE_IPCQ component (ADR-0023 D4 weak fairness)."""
from __future__ import annotations

import numpy as np

from kernbench.ccl.testing import run_kernel_in_mock


def kernel_round_robin(t_ptr, n_elem, tl):
    """Each PE sends one tile E then receives N-1 tiles via round-robin.
    Uses TensorHandle math (PE_MATH) so Phase 2 produces correct HBM
    contents under SimPy + op_log replay."""
    rank = tl.program_id(axis=0)
    world_size = tl.num_programs(axis=0)
    nbytes = n_elem * 2

    pe_addr = t_ptr + rank * nbytes
    acc = tl.load(pe_addr, shape=(n_elem,), dtype="f16")
    current = acc

    for _step in range(world_size - 1):
        tl.send(dir="E", src=current)
        # No direction → round-robin
        recv = tl.recv(shape=(n_elem,), dtype="f16")
        acc = acc + recv
        current = recv  # forward W's tile to E next round

    tl.store(pe_addr, acc)


def test_round_robin_recv_mock_runtime():
    n_elem = 8
    inputs = [
        np.full((n_elem,), float(r + 1), dtype=np.float16)
        for r in range(4)
    ]
    expected = sum(inputs)  # [10,...]

    outputs = run_kernel_in_mock(
        kernel_fn=kernel_round_robin,
        world_size=4,
        topology="ring_1d",
        inputs=inputs,
        kernel_args=(n_elem,),
    )
    for r in range(4):
        assert np.allclose(outputs[r], expected)