"""milestone-1h-gqa decode: Case-6 composite-command study.

Three command-form variants of the Case-6 (Cube-SP × PE-SP) long-context
decode kernel, swept over S_kv so the local-attention tile loop runs
multiple tiles (S_local = S_kv/(C·P) > TILE_S_KV = 1024):

    primitive           tl.dot + primitive online-softmax merge (baseline).
    composite           per-tile GEMMs as tl.composite GEMM commands.
    composite_extended  per-tile attention as a Q·Kᵀ composite + a
                        softmax_merge recipe composite (ADR-0065).

The placement, DP policy, and (m, ℓ, O) reduce are identical across
variants; only the per-tile command form differs. The sweep writes
per-(variant, S_kv) latency, engine occupancy, and PE_CPU command count
to sweep_decode_composite.json so the comparative plot
(paper_plot_gqa_decode_long_ctx_composite.py) can read off the
CPU-offload win.

Run in op_log mode (enable_data=False): latency / dispatch only — recipe
data-mode numeric parity is a separate follow-up (DDD-0065).
"""
from __future__ import annotations

import json
from pathlib import Path

from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import (  # noqa: E501
    gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext import (  # noqa: E501
    gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import (  # noqa: E501
    gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel,
)
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
from kernbench.policy.placement.dp import DPPolicy

_OUTPUT_DIR = (
    Path(__file__).resolve().parents[2]
    / "1H_milestone_output" / "gqa" / "long_ctx"
)
_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json"


# ── Variant + S_kv registry ──────────────────────────────────────────


# Each kernel implements the same Case-6 placement; only the per-tile
# command form differs (see module docstring).
_VARIANT_KERNELS = {
    "primitive_tiled": gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel,
    "composite": gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
    "composite_extended":
        gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
}
_VARIANTS = ("primitive_tiled", "composite", "composite_extended")

# Op-count (PE_CPU dispatch) is computed at emit time (exact, instant), so
# it spans up to the 1M production point (S_local = 1M/64 = 16384, 16
# tiles): the per-PE work is Q·Kᵀ (8,128)·(128,16384) and P·V
# (8,16384)·(16384,128). End-to-end latency needs the data-mode engine,
# whose cost scales with S_kv, so it is swept only over a tractable range.
_S_KV_OPCOUNT = (8192, 65_536, 131_072, 262_144, 524_288, 1_048_576)
_S_KV_LATENCY = (8192, 32_768, 65_536, 131_072)

# LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs), one decode step.
_BASE_PARAMS = dict(C=8, P=8, T_q=1, d_head=128, h_q=8, h_kv=1)


# ── Per-(variant, S_kv) runner ───────────────────────────────────────


def _run_panel_fn(variant: str, S_kv: int):
    kernel = _VARIANT_KERNELS[variant]
    p = _BASE_PARAMS
    panel = f"decode_long_ctx_composite_{variant}_s{S_kv}"

    def _bench_fn(ctx):
        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
        dp_full = DPPolicy(cube="replicate", pe="replicate",
                           num_cubes=p["C"], num_pes=p["P"])
        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
                         num_cubes=p["C"], num_pes=p["P"])
        q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
                      dtype="f16", dp=dp_full, name=f"{panel}_q")
        k = ctx.zeros((S_kv, p["h_kv"] * p["d_head"]),
                      dtype="f16", dp=dp_kv, name=f"{panel}_k")
        v = ctx.zeros((S_kv, p["h_kv"] * p["d_head"]),
                      dtype="f16", dp=dp_kv, name=f"{panel}_v")
        o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
                      dtype="f16", dp=dp_full, name=f"{panel}_o")
        ctx.launch(panel, kernel, q, k, v, o,
                   p["T_q"], S_kv, p["h_q"], p["h_kv"], p["d_head"],
                   p["C"], p["P"], _auto_dim_remap=False)

    return _bench_fn


def _end_to_end_ns(op_log) -> float:
    if not op_log:
        return 0.0
    return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)


def _emit_dispatch(variant: str, S_kv: int) -> tuple[int, float]:
    """PE_CPU dispatch (# commands, summed cycles) the kernel emits at the
    lrab center rank (cube 6, pe 0) — computed at command-emit time, so it
    is exact and S_kv-independent for the composite forms (no engine)."""
    from kernbench.common.pe_commands import PeCpuOverheadCmd
    from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
    from kernbench.triton_emu.tl_context import TLContext, run_kernel

    p = _BASE_PARAMS
    tl = TLContext(
        pe_id=0, num_programs=p["P"], cost_model=DEFAULT_PE_COST_MODEL,
        cube_id=6, num_cubes=p["C"], scratch_base=1 << 61, scratch_size=1 << 20,
    )
    run_kernel(
        _VARIANT_KERNELS[variant], tl,
        0x1000, 0x2000, 0x3000, 0x4000,
        p["T_q"], S_kv, p["h_q"], p["h_kv"], p["d_head"], p["C"], p["P"],
    )
    cmds = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
    return len(cmds), sum(c.cycles for c in cmds)


def _engine_latency_ns(variant: str, S_kv: int, topology: str) -> float:
    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=_run_panel_fn(variant, S_kv),
        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"gqa-decode-composite {variant}@{S_kv} failed: {result.completion}"
        )
    return _end_to_end_ns(result.engine.op_log)


# ── Sweep entry (called by the umbrella milestone_1h_gqa) ────────────


def run_sweep(topology: str = "topology.yaml") -> int:
    """Emit-level dispatch over the full S_kv range (to 1M) + engine
    latency over the tractable range; write sweep.json. Returns row count."""
    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    rows = []
    for S_kv in _S_KV_OPCOUNT:
        for variant in _VARIANTS:
            n_cmds, cycles = _emit_dispatch(variant, S_kv)
            latency = (
                _engine_latency_ns(variant, S_kv, topology)
                if S_kv in _S_KV_LATENCY else None
            )
            rows.append({
                "variant": variant,
                "S_kv": S_kv,
                **_BASE_PARAMS,
                "pe_cpu_cmd_count": n_cmds,
                "pe_cpu_dispatch_cycles": cycles,
                "latency_ns": latency,
            })
    sweep = {
        "version": 2,
        "variants": list(_VARIANTS),
        "s_kv_opcount": list(_S_KV_OPCOUNT),
        "s_kv_latency": list(_S_KV_LATENCY),
        "rows": rows,
    }
    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
    print(f"  gqa-decode-composite: {len(rows)} rows -> {_SWEEP_JSON}")
    return len(rows)
