"""Generate a markdown summary for attention path benchmark JSON output."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def ms(value: float) -> float:
    return value * 1000.0


def ok_results(data: dict) -> list[dict]:
    return [item for item in data["results"] if item.get("status") == "ok"]


def timing_ms(result: dict, name: str, field: str = "mean_s") -> float:
    return ms(result["timings"][name][field])


def label(result: dict) -> str:
    layout = result.get("layout_mode", "")
    return f"{result['backend']}:{layout}" if layout else result["backend"]


def timing_names(result: dict) -> tuple[str, ...]:
    layout = result.get("layout_mode")
    if layout == "sol_fused":
        return ("qkv_to_bshd", "attention_kernel", "output_reshape")
    if layout == "sol_native":
        return ("q_bshd_contiguous", "k_bshd_contiguous", "v_bshd_contiguous", "attention_kernel", "output_reshape")
    return ("q_transpose_contiguous", "k_transpose_contiguous", "v_transpose_contiguous", "attention_kernel", "output_reshape")


def table(results: list[dict]) -> list[str]:
    lines = ["| Segment | Backend | Kernel mean ms | Total mean ms | Layout mean ms | Kernel p50 ms | Total p50 ms |", "| --- | --- | ---: | ---: | ---: | ---: | ---: |"]
    for item in sorted(results, key=lambda row: (row["segment"], row["layout_attention_total_mean_s"])):
        total = ms(item["layout_attention_total_mean_s"])
        kernel = timing_ms(item, "attention_kernel")
        layout = total - kernel
        p50_total = sum(timing_ms(item, name, "p50_s") for name in timing_names(item))
        lines.append(f"| {item['segment']} | {label(item)} | {kernel:.3f} | {total:.3f} | {layout:.3f} | {timing_ms(item, 'attention_kernel', 'p50_s'):.3f} | {p50_total:.3f} |")
    return lines


def fastest_by_segment(results: list[dict]) -> list[str]:
    segments = sorted({item["segment"] for item in results})
    lines = []
    for segment in segments:
        segment_results = [item for item in results if item["segment"] == segment]
        best = min(segment_results, key=lambda item: item["layout_attention_total_mean_s"])
        best_kernel = min(segment_results, key=lambda item: item["timings"]["attention_kernel"]["mean_s"])
        lines.append(f"- `{segment}` fastest total: `{label(best)}` at {ms(best['layout_attention_total_mean_s']):.3f} ms; fastest kernel: `{label(best_kernel)}` at {timing_ms(best_kernel, 'attention_kernel'):.3f} ms.")
    return lines


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--whole", type=Path, required=True)
    parser.add_argument("--segments", type=Path, required=True)
    parser.add_argument("--extra", type=Path, action="append", default=[])
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    whole = json.loads(args.whole.read_text(encoding="utf-8"))
    segments = json.loads(args.segments.read_text(encoding="utf-8"))
    extras = [json.loads(path.read_text(encoding="utf-8")) for path in args.extra]
    whole_ok = ok_results(whole)
    segment_ok = ok_results(segments)
    extra_ok = [item for data in extras for item in ok_results(data)]
    extra_lines = table(extra_ok) if extra_ok else ["No extra runs."]
    failures = [item for item in segments["results"] + whole["results"] + [item for data in extras for item in data["results"]] if item.get("status") != "ok"]

    lines = [
        "# H3 Attention Path Profile",
        "",
        f"Whole-sequence input shape: `{whole['metadata']['h_msa_shape']}`; segments: `{whole['segments']}`.",
        f"Measurements: `{whole['warmup']}` warmups and `{whole['iterations']}` iterations per backend.",
        "",
        "## Key Findings",
        "",
        *fastest_by_segment(whole_ok + segment_ok + extra_ok),
        "- Sol-native BSHD removes the generic HND round trip and keeps exact Sol output parity.",
        "- The current fused QKV split/layout prototype is opt-in because it is not faster than three PyTorch contiguous copies yet.",
        "- Full-sequence Sol-Attn stays under the initial 50-55 ms target for attention plus layout on this captured H3 shape.",
        "- Q/K/V layout remains the next memory-bandwidth target; the Sage-style output reshape is not worth attacking for Sol.",
        "- `sdpa` is fastest for tiny text/secondary segments but loses heavily on the video-dominated path, so it is only interesting for segment-specialized dispatch.",
        "",
        "## Whole Sequence",
        "",
        *table(whole_ok),
        "",
        "## Extra Whole-Sequence Layout Runs",
        "",
        *extra_lines,
        "",
        "## Segment Sweep",
        "",
        *table(segment_ok),
    ]
    if failures:
        lines.extend(["", "## Expected Failures", ""])
        lines.extend(f"- `{item['segment']}` / `{label(item)}`: `{item['error']}`" for item in failures)
    lines.extend(["", "## Source Files", "", f"- Whole JSON: `{args.whole}`", f"- Segment JSON: `{args.segments}`"])
    lines.extend(f"- Extra JSON: `{path}`" for path in args.extra)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text("\n".join(lines) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()