"""Phase 1 spec test for the 2D row-then-col AllReduce-mlo decode kernel.

The 2D kernel decomposes a ``(mesh_rows × mesh_cols)`` cube sub-mesh into a
two-stage AllReduce-mlo: stage 1 reduces across columns within each row
(E/W edges), stage 2 reduces across rows within each column (N/S edges).
After both stages every cube holds the same final ``(m, ℓ, o)``.

This module is Phase 1 of C2 (see CLAUDE.md change protocol): it pins
the kernel's interface and observable behavior. Production code for the
kernel lands in Phase 2; until then this file fails to import.

Test shapes (run on default ``topology.yaml`` — 2 SIPs × 4×4 cube_mesh):

  1×4 sub-mesh (4 cubes, row 0 only)
     Degenerates to a row-only AllReduce — equivalent in step count to
     the existing 1D kernel at n_ranks=4. Verifies the kernel reduces
     correctly when mesh_rows=1 (stage 2 collapses to no-op).

  2×4 sub-mesh (8 cubes, rows 0+1)
     The 8-KV-group target. Verifies that cubes 4..7 use ``dir="N"``
     (not ``dir="W"``) to reach row 0 — surfacing the IpcqInvalidDirection
     bug that the 1D kernel hit at cube 4 (rank 4, no W neighbor).

  4×4 sub-mesh (16 cubes, full SIP)
     Full-SIP scale. Verifies the algorithm fans out over (cols-1)=3
     row steps followed by (rows-1)=3 col steps.
"""
from __future__ import annotations

from pathlib import Path

from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
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[2] / "topology.yaml"

S_Q = 1
S_KV_PER_RANK = 16
H_Q = 1
H_KV = 1
D_HEAD = 64
DTYPE = "f16"


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 _run_2d(mesh_rows: int, mesh_cols: int, cube_start: int = 0):
    """Build a bench_fn and run it on the default topology."""
    n_cubes = mesh_rows * mesh_cols
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
        dp_full = DPPolicy(cube="replicate", pe="replicate",
                           num_cubes=n_cubes, num_pes=8,
                           cube_start=cube_start)
        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
                         num_cubes=n_cubes, num_pes=8,
                         cube_start=cube_start)
        q = ctx.zeros((S_Q, H_Q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name="q")
        k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name="k")
        v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name="v")
        o = ctx.empty((S_Q, H_Q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name="o")
        ctx.launch(
            f"mesh_mlo_2d_{mesh_rows}x{mesh_cols}_start{cube_start}",
            attention_mesh_mlo_2d_kernel,
            q, k, v, o,
            S_Q, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
            mesh_rows, mesh_cols,
            1,  # rank_axis=1 → cube-level ring
            cube_start,
            _auto_dim_remap=False,
        )

    return run_bench(
        topology=topo,
        bench_fn=_bench_fn,
        device=resolve_device(None),
        engine_factory=_engine_factory,
    )


def test_2d_kernel_1x4_row_only():
    """1×4 sub-mesh: row-only AllReduce, stage 2 collapses to no-op."""
    result = _run_2d(mesh_rows=1, mesh_cols=4)
    assert result.completion.ok, (
        f"1x4: completion not ok - {result.completion}"
    )


def test_2d_kernel_2x4_eight_cubes():
    """2×4 sub-mesh: the 8-KV-group target.

    Verifies that cube 4 (row 1, col 0) uses ``dir="N"`` to reach cube 0
    (row 0, col 0) for stage 2, not ``dir="W"`` — the 1D kernel hit
    IpcqInvalidDirection here.
    """
    result = _run_2d(mesh_rows=2, mesh_cols=4)
    assert result.completion.ok, (
        f"2x4: completion not ok - {result.completion}"
    )


def test_2d_kernel_4x4_full_sip():
    """4×4 sub-mesh: full-SIP scale (16 cubes)."""
    result = _run_2d(mesh_rows=4, mesh_cols=4)
    assert result.completion.ok, (
        f"4x4: completion not ok - {result.completion}"
    )


def test_2d_kernel_2x4_at_cube_start_eight():
    """2×4 sub-mesh at cube_start=8 (cubes 8..15, rows 2..3).

    The second KV-group per SIP in the 8-KV-group Llama-70B headline.
    Verifies the kernel converts ``program_id(axis=1)`` (physical cube
    id) back to launch-local rank via ``cube_start`` — without that
    subtraction, cube 8 would compute my_row=2 (out of sub-mesh bounds)
    and deadlock waiting on cube 4 which isn't in the launch.
    """
    result = _run_2d(mesh_rows=2, mesh_cols=4, cube_start=8)
    assert result.completion.ok, (
        f"2x4 @ cube_start=8: completion not ok - {result.completion}"
    )
