"""milestone-1h-gqa: umbrella GQA bench for the 1H code-sign milestone.

Single ``@bench`` entry that drives all GQA panels by delegating to
internal helpers under ``kernbench.benches.gqa_helpers``. Mirrors the
``milestone_1h_gemm`` umbrella pattern.

Currently exercises (long-context only — short-context panels are
future work):
  - 4-cases prefill comparative study (gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases)
  - 4-cases decode comparative study  (gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases)
  - Case-6 composite-command study    (gqa_helpers.long_ctx.gqa_decode_long_ctx_composite)

Each sub-sweep writes its own ``sweep_*.json`` into the shared output
dir ``benches/1H_milestone_output/gqa/gqa_long_ctx/``.
Selection via the env var ``GQA_1H_SWEEPS=prefill,decode`` (default
runs prefill+decode). The composite study is opt-in (it sweeps the
data-mode engine over several S_kv points): ``GQA_1H_SWEEPS=composite``.

Gated by ``GQA_1H_RUN=1`` to keep CI fast.
"""
from __future__ import annotations

import os

from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
    run_sweep as _run_decode_sweep,
)
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
    run_sweep as _run_composite_sweep,
)
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
    run_sweep as _run_prefill_cb_sweep,
)
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases import (
    run_sweep as _run_prefill_sweep,
)
from kernbench.benches.registry import bench
from kernbench.policy.placement.dp import DPPolicy


@bench(
    name="milestone-1h-gqa",
    description=(
        "Umbrella GQA milestone — drives long-context prefill + decode "
        "4-cases sweeps on the LLaMA-3.1-70B single-KV-head group "
        "(8 cubes × 8 PEs). Gated by GQA_1H_RUN=1."
    ),
)
def run(torch) -> None:
    """Drive selected GQA sub-sweeps; each writes its own sweep_*.json.

    Env vars:
      GQA_1H_RUN=1                    (required gate)
      GQA_1H_TOPOLOGY=topology.yaml   (override topology path)
      GQA_1H_SWEEPS=prefill,decode    (default: both; comma-separated)
    """
    if not os.environ.get("GQA_1H_RUN"):
        raise RuntimeError("milestone-1h-gqa needs GQA_1H_RUN=1.")

    topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
    requested = os.environ.get("GQA_1H_SWEEPS", "prefill,decode")
    sweeps = [s.strip() for s in requested.split(",") if s.strip()]

    runners = {
        "prefill": _run_prefill_sweep,
        "decode": _run_decode_sweep,
        "composite": _run_composite_sweep,
        "prefill_cb": _run_prefill_cb_sweep,
    }
    unknown = [s for s in sweeps if s not in runners]
    if unknown:
        raise RuntimeError(
            f"GQA_1H_SWEEPS contains unknown sweep(s) {unknown}; "
            f"valid: {sorted(runners)}"
        )

    for s in sweeps:
        runners[s](topology)

    # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out) — the
    # sub-sweeps each spin up their own GraphEngine via ``run_bench``,
    # so this outer @bench needs to submit at least one request.
    torch.zeros(
        (1, 1), dtype="f16",
        dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
        name="milestone_1h_gqa_sentinel",
    )