"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV.

Wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels through 4
panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing
per-panel ``op_log_summary`` into ``sweep.json``. Independent from the
existing ``milestone-gqa-llama70b`` validation-scale bench (which stays
on the legacy baseline kernels).

Restrictions:
  - Decode side capped at C ≤ 4 (1D chain reduce; 2D mesh wiring on
    the decode panels deferred to the 4-cases decode comparative study)
  - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
    headline deferred per ADR-0060 §B-item-1)
  - No figure renderers (defer to a separate cycle)

Panels:
  single_user_prefill_gqa         : prefill C=1, T_q=4, S_kv=16
  multi_user_prefill_gqa          : prefill C=4 Ring KV, T_q=4, S_kv=16
  single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
                                     intra-CUBE PE-SP (all 64 ranks),
                                     T_q=S_kv=1K, d_head=128 — the
                                     LLaMA-3.1-70B single-KV-group target
                                     (scratch-limited; 32K headline awaits
                                     Q-axis kernel tiling)
  single_user_decode_gqa          : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64
                                    (M-fold + intra-cube row-then-col chain)
  multi_user_decode_gqa           : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128
                                    (M-fold + 2-level chain reduce-to-root)

Gated by ``GQA_HEADLINE_RUN=1``.
"""
from __future__ import annotations

import json
import os
from pathlib import Path

from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel
from kernbench.benches.registry import bench
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import (
    configure_sfr_intercube_multisip,
    configure_sfr_intercube_ring,
)
from kernbench.policy.placement.dp import DPPolicy

_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline"
_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"

# ── Panel configs ────────────────────────────────────────────────────

_DTYPE = "f16"
_D_HEAD = 64
_T_Q_PREFILL = 4
_T_Q_DECODE = 1
_S_KV_PREFILL = 16
_H_Q_DECODE = 8        # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8
_H_KV_DECODE = 1

_PANELS = (
    "single_user_prefill_gqa",
    "multi_user_prefill_gqa",
    "single_kv_group_prefill_gqa_c8_p8",
    "single_user_decode_gqa",
    "multi_user_decode_gqa",
)

# Each entry: (kind, panel-specific params)
_PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
    "single_user_prefill_gqa": ("prefill", {"C": 1, "S_kv": _S_KV_PREFILL}),
    "multi_user_prefill_gqa":  ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}),
    "single_kv_group_prefill_gqa_c8_p8": ("prefill", {
        "C": 8, "P": 8,
        # T_q = S_kv = 1024 (one-shot long-context prefill, scratch-limited).
        # The bootstrap section of the prefill kernel leaves K_t/V_t/scores/
        # exp_scores persistent (outside tl.scratch_scope), inflating the
        # baseline. At T_q=2048 the peak hits ~1.05 MB vs 1.0 MB budget; at
        # 1024 it fits comfortably. The true LLaMA 32K headline awaits a
        # future increment to (a) add Q-axis tiling and (b) move the
        # bootstrap into scratch discipline.
        "T_q": 1_024, "S_kv": 1_024,
        "d_head": 128,
    }),
    "single_user_decode_gqa":  ("decode",  {"C": 1, "P": 8, "S_kv": 64}),
    "multi_user_decode_gqa":   ("decode",  {"C": 4, "P": 8, "S_kv": 128}),
}


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


# ── Per-kind launch helpers ──────────────────────────────────────────


def _run_prefill_panel(
    ctx, *, panel: str, C: int, S_kv: int,
    P: int = 1,
    T_q: int = _T_Q_PREFILL,
    d_head: int = _D_HEAD,
) -> None:
    if C > 1:
        mesh_w = int(ctx.spec["sip"]["cube_mesh"]["w"])
        if C <= mesh_w:
            configure_sfr_intercube_ring(
                ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
            )
        else:
            if C % mesh_w != 0:
                raise ValueError(
                    f"C={C} > mesh_w={mesh_w} requires C divisible by mesh_w"
                )
            configure_sfr_intercube_ring(
                ctx.engine, ctx.spec, _ccl_cfg(),
                submesh_shape=(C // mesh_w, mesh_w),
            )
    # Q and O switch to pe="row_wise" when intra-CUBE PE-SP is active
    # (ADR-0060 §5.5 + §B-item-3: disjoint query-row split across PEs).
    q_pe = "row_wise" if P > 1 else "replicate"
    o_pe = "row_wise" if P > 1 else "replicate"
    dp_q = DPPolicy(cube="replicate", pe=q_pe,
                    num_cubes=C, num_pes=P)
    dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
                     pe="replicate", num_cubes=C, num_pes=P)
    dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
                    pe=o_pe, num_cubes=C, num_pes=P)
    q = ctx.zeros((T_q, d_head),
                  dtype=_DTYPE, dp=dp_q, name=f"{panel}_q")
    k = ctx.zeros((S_kv, d_head),
                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, d_head),
                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((T_q * C, d_head),
                  dtype=_DTYPE, dp=dp_o, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_prefill_long_kernel,
        q, k, v, o,
        T_q, S_kv, d_head, C, P,
        _auto_dim_remap=False,
    )


def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None:
    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
    dp_full = DPPolicy(cube="replicate", pe="replicate",
                       num_cubes=C, num_pes=P)
    dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
                     pe="row_wise", num_cubes=C, num_pes=P)
    q = ctx.zeros((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
                  dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
    k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
                  dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_decode_long_kernel,
        q, k, v, o,
        _T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P,
        _auto_dim_remap=False,
    )


def _make_bench_fn(panel: str):
    kind, params = _PANEL_DISPATCH[panel]

    def _bench_fn(ctx):
        if kind == "prefill":
            _run_prefill_panel(ctx, panel=panel, **params)
        else:
            _run_decode_panel(ctx, panel=panel, **params)

    return _bench_fn


# ── Op-log summary ──────────────────────────────────────────────────


def _summarize_op_log(op_log) -> dict[str, int]:
    """Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write)."""
    gemm_count = 0
    ipcq_copy_count = 0
    dma_read_count = 0
    dma_write_count = 0
    for r in op_log:
        if r.op_kind == "gemm":
            gemm_count += 1
        elif r.op_name == "dma_read":
            dma_read_count += 1
        elif r.op_name == "dma_write":
            dma_write_count += 1
        elif r.op_name == "ipcq_copy":
            ipcq_copy_count += 1
    return {
        "gemm_count": gemm_count,
        "ipcq_copy_count": ipcq_copy_count,
        "dma_read_count": dma_read_count,
        "dma_write_count": dma_write_count,
    }


def _run_panel(panel: str, topology: str) -> dict:
    """Run one panel in a fresh engine; return its row dict."""
    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

    topo = resolve_topology(topology)
    result = run_bench(
        topology=topo, bench_fn=_make_bench_fn(panel),
        device=resolve_device(None),
        engine_factory=lambda t, d: GraphEngine(
            getattr(t, "topology_obj", t), enable_data=True,
        ),
    )
    if not result.completion.ok:
        raise RuntimeError(
            f"milestone-gqa-headline panel {panel!r} failed: "
            f"{result.completion}"
        )
    kind, params = _PANEL_DISPATCH[panel]
    return {
        "panel": panel,
        "kind": kind,
        **params,
        "op_log_summary": _summarize_op_log(result.engine.op_log),
    }


# ── Bench entry ──────────────────────────────────────────────────────


@bench(
    name="milestone-gqa-headline",
    description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).",
)
def run(torch) -> None:
    """Drive 4 headline panels through the new GQA kernels; write sweep.json.

    Gated by GQA_HEADLINE_RUN=1.
    """
    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    if not os.environ.get("GQA_HEADLINE_RUN"):
        raise RuntimeError(
            "milestone-gqa-headline needs GQA_HEADLINE_RUN=1."
        )

    topology = os.environ.get("GQA_HEADLINE_TOPOLOGY", "topology.yaml")
    rows = [_run_panel(panel, topology) for panel in _PANELS]
    sweep = {
        "version": 1,
        "panels": list(_PANELS),
        "config": {
            "T_q_prefill": _T_Q_PREFILL,
            "T_q_decode": _T_Q_DECODE,
            "S_kv_prefill": _S_KV_PREFILL,
            "h_q_decode": _H_Q_DECODE,
            "h_kv_decode": _H_KV_DECODE,
            "d_head": _D_HEAD,
        },
        "rows": rows,
    }
    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
    print(f"  milestone-gqa-headline: {len(rows)} rows -> {_SWEEP_JSON}")

    # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out).
    torch.zeros(
        (1, 1), dtype="f16",
        dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
        name="milestone_gqa_headline_sentinel",
    )