#!/usr/bin/env python3
"""Summarize existing benchmark artifacts for characterization review.

This is a CPU-only companion to ``analyze.py``. It does not run benchmarks.
It reads completed output directories and produces an audit-oriented package
that helps decide which TODO claims are currently supported by existing data
and which still need fresh GPU runs or additional instrumentation.
"""

from __future__ import annotations

import argparse
import csv
import datetime as dt
import json
import math
import statistics
import subprocess
from pathlib import Path
from typing import Any


JsonDict = dict[str, Any]


DEFAULT_RUNS = [
    "outputs/gpu_ab_combined",
    "outputs/gpu_ab_pdsep",
    "outputs/contention_16s_ts10",
    "outputs/contention_16s_elastic",
    "outputs/combined_1000req",
    "outputs/exp3_pd_sep_tp1_mooncake",
]


def main() -> None:
    args = parse_args()
    out_dir = args.output_dir
    out_dir.mkdir(parents=True, exist_ok=True)

    run_dirs = [Path(p) for p in (args.runs or DEFAULT_RUNS)]
    summaries = [summarize_run(path) for path in run_dirs]
    comparisons = build_comparisons(summaries)
    claim_matrix = build_claim_matrix(summaries, comparisons)
    risk_register = build_risk_register(summaries)

    write_json(out_dir / "run_summaries.json", summaries)
    write_json(out_dir / "comparisons.json", comparisons)
    write_json(out_dir / "claim_matrix.json", claim_matrix)
    write_json(out_dir / "reviewer_risk_register.json", risk_register)
    (out_dir / "current_results.md").write_text(
        render_current_results(summaries, comparisons, claim_matrix, risk_register),
        encoding="utf-8",
    )
    (out_dir / "characterization_claim_matrix.md").write_text(
        render_claim_matrix(claim_matrix),
        encoding="utf-8",
    )
    (out_dir / "reviewer_risk_register.md").write_text(
        render_risk_register(risk_register),
        encoding="utf-8",
    )
    (out_dir / "all_figures_index.md").write_text(
        render_figures_index(summaries),
        encoding="utf-8",
    )
    (out_dir / "reproduction_commands.sh").write_text(
        render_reproduction_commands(args, run_dirs),
        encoding="utf-8",
    )
    print(f"Wrote run summary package to {out_dir}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Summarize existing characterization-relevant output dirs.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--runs",
        nargs="*",
        default=None,
        help="Output directories to summarize. Defaults to a small curated set.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("analysis/characterization/current_results"),
        help="Directory for generated review artifacts.",
    )
    return parser.parse_args()


def summarize_run(path: Path) -> JsonDict:
    metrics_summary = load_json(path / "metrics.summary.json")
    metrics_rows = load_jsonl(path / "metrics.jsonl")
    gpu_summary = summarize_gpu(path / "gpu_util.csv")
    breakdown_summary = summarize_breakdown(path / "breakdown.json")
    apc_summary = summarize_apc(path / "apc.txt")
    return {
        "run": str(path),
        "exists": path.exists(),
        "metrics_summary_available": bool(metrics_summary),
        "metrics_jsonl_rows": len(metrics_rows),
        "request_count": first_present(metrics_summary, ["request_count"]),
        "success_count": first_present(metrics_summary, ["success_count"]),
        "error_count": first_present(metrics_summary, ["error_count"]),
        "wall_clock_s": first_present(metrics_summary, ["wall_clock_s"]),
        "latency_stats_s": metrics_summary.get("latency_stats_s"),
        "ttft_stats_s": metrics_summary.get("ttft_stats_s"),
        "tpot_stats_s": metrics_summary.get("tpot_stats_s"),
        "prefix_cache_hit_ratio": metrics_summary.get("prefix_cache_hit_ratio"),
        "external_cache_hit_ratio": metrics_summary.get("external_cache_hit_ratio"),
        "session_summary": summarize_sessions(metrics_rows),
        "gpu_summary": gpu_summary,
        "breakdown_summary": breakdown_summary,
        "apc_summary": apc_summary,
        "artifact_availability": {
            "metrics_summary_json": (path / "metrics.summary.json").exists(),
            "metrics_jsonl": (path / "metrics.jsonl").exists(),
            "gpu_util_csv": (path / "gpu_util.csv").exists(),
            "breakdown_json": (path / "breakdown.json").exists(),
            "apc_txt": (path / "apc.txt").exists(),
        },
    }


def summarize_sessions(rows: list[JsonDict]) -> JsonDict:
    if not rows:
        return {
            "status": "unavailable",
            "reason": "metrics.jsonl missing",
        }
    sessions: dict[str, JsonDict] = {}
    input_values = []
    output_values = []
    cached_values = []
    for row in rows:
        sid = str(row.get("session_id", ""))
        item = sessions.setdefault(
            sid,
            {
                "turns": 0,
                "input_tokens": 0.0,
                "output_tokens": 0.0,
                "cached_tokens": 0.0,
            },
        )
        inp = to_float(row.get("input_length")) or 0.0
        out = to_float(row.get("actual_output_tokens")) or to_float(row.get("output_length")) or 0.0
        cached = to_float(row.get("cached_tokens")) or 0.0
        item["turns"] += 1
        item["input_tokens"] += inp
        item["output_tokens"] += out
        item["cached_tokens"] += cached
        input_values.append(inp)
        output_values.append(out)
        cached_values.append(cached)
    per_session_input = [v["input_tokens"] for v in sessions.values()]
    return {
        "status": "available",
        "request_input_tokens": stats(input_values),
        "request_output_tokens": stats(output_values),
        "request_cached_tokens": stats(cached_values),
        "session_count": len(sessions),
        "turns_per_session": stats([v["turns"] for v in sessions.values()]),
        "session_input_tokens": stats(per_session_input),
        "top_session_input_fraction": top_contribution(per_session_input),
    }


def summarize_gpu(path: Path) -> JsonDict:
    if not path.exists():
        return {
            "status": "unavailable",
            "reason": "gpu_util.csv missing",
        }
    values: dict[str, list[float]] = {}
    with path.open() as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            gpu = str(row.get("gpu", ""))
            util = to_float(row.get("util_pct"))
            if gpu and util is not None:
                values.setdefault(gpu, []).append(util)
    means = {gpu: statistics.fmean(vals) for gpu, vals in values.items() if vals}
    if not means:
        return {
            "status": "unavailable",
            "reason": "gpu_util.csv had no util_pct rows",
        }
    mean_values = list(means.values())
    return {
        "status": "available",
        "gpu_count": len(means),
        "per_gpu_mean_util_pct": means,
        "mean_util_pct": statistics.fmean(mean_values),
        "stddev_across_gpu_mean_util_pct": statistics.pstdev(mean_values),
        "max_mean_util_pct": max(mean_values),
        "min_mean_util_pct": min(mean_values),
        "max_min_ratio": max(mean_values) / max(min(mean_values), 1e-9),
    }


def summarize_breakdown(path: Path) -> JsonDict:
    if not path.exists():
        return {
            "status": "unavailable",
            "reason": "breakdown.json missing",
        }
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        return {
            "status": "unavailable",
            "reason": f"failed to parse breakdown: {exc}",
        }
    rows: list[JsonDict]
    if isinstance(data, list):
        rows = [r for r in data if isinstance(r, dict)]
    elif isinstance(data, dict):
        rows = data.get("records") if isinstance(data.get("records"), list) else []
        if not rows:
            rows = [data]
    else:
        rows = []
    mode_counts: dict[str, int] = {}
    route_counts: dict[str, int] = {}
    for row in rows:
        mode = row.get("mode") or row.get("execution_mode") or row.get("route_class")
        route = row.get("route") or row.get("decision") or row.get("policy")
        if mode is not None:
            mode_counts[str(mode)] = mode_counts.get(str(mode), 0) + 1
        if route is not None:
            route_counts[str(route)] = route_counts.get(str(route), 0) + 1
    return {
        "status": "available",
        "row_count": len(rows),
        "mode_counts": mode_counts,
        "route_counts": route_counts,
        "field_sample": sorted(rows[0].keys()) if rows else [],
    }


def summarize_apc(path: Path) -> JsonDict:
    if not path.exists():
        return {
            "status": "unavailable",
            "reason": "apc.txt missing",
        }
    text = path.read_text(encoding="utf-8", errors="replace")
    return {
        "status": "available",
        "line_count": len(text.splitlines()),
        "preview": "\n".join(text.splitlines()[:20]),
    }


def build_comparisons(summaries: list[JsonDict]) -> list[JsonDict]:
    by_run = {s["run"]: s for s in summaries}
    pairs = [
        ("combined_vs_pdsep_200", "outputs/gpu_ab_combined", "outputs/gpu_ab_pdsep"),
        ("contention_baseline_vs_elastic_500", "outputs/contention_16s_ts10", "outputs/contention_16s_elastic"),
        ("combined_1000_vs_pdsep_mooncake", "outputs/combined_1000req", "outputs/exp3_pd_sep_tp1_mooncake"),
    ]
    out = []
    for name, base, variant in pairs:
        if base not in by_run or variant not in by_run:
            continue
        out.append(compare_pair(name, by_run[base], by_run[variant]))
    return out


def compare_pair(name: str, base: JsonDict, variant: JsonDict) -> JsonDict:
    return {
        "name": name,
        "baseline": base["run"],
        "variant": variant["run"],
        "request_count": [base.get("request_count"), variant.get("request_count")],
        "success_count": [base.get("success_count"), variant.get("success_count")],
        "error_count": [base.get("error_count"), variant.get("error_count")],
        "ttft_p50_delta_pct": pct_delta(stat_value(base, "ttft_stats_s", "p50"), stat_value(variant, "ttft_stats_s", "p50")),
        "ttft_p90_delta_pct": pct_delta(stat_value(base, "ttft_stats_s", "p90"), stat_value(variant, "ttft_stats_s", "p90")),
        "e2e_p50_delta_pct": pct_delta(stat_value(base, "latency_stats_s", "p50"), stat_value(variant, "latency_stats_s", "p50")),
        "e2e_p90_delta_pct": pct_delta(stat_value(base, "latency_stats_s", "p90"), stat_value(variant, "latency_stats_s", "p90")),
        "tpot_p90_delta_pct": pct_delta(stat_value(base, "tpot_stats_s", "p90"), stat_value(variant, "tpot_stats_s", "p90")),
        "wall_clock_delta_pct": pct_delta(base.get("wall_clock_s"), variant.get("wall_clock_s")),
        "gpu_mean_util": [
            nested(base, ["gpu_summary", "mean_util_pct"]),
            nested(variant, ["gpu_summary", "mean_util_pct"]),
        ],
        "gpu_imbalance_ratio": [
            nested(base, ["gpu_summary", "max_min_ratio"]),
            nested(variant, ["gpu_summary", "max_min_ratio"]),
        ],
    }


def build_claim_matrix(summaries: list[JsonDict], comparisons: list[JsonDict]) -> list[JsonDict]:
    has_gpu = any((s.get("gpu_summary") or {}).get("status") == "available" for s in summaries)
    has_metrics = any(s.get("metrics_summary_available") for s in summaries)
    return [
        {
            "claim": "Batch 0 substrate audit is only partially complete for existing runs.",
            "status": "partially_supported",
            "supporting_data": "metrics.jsonl lacks actual dispatch/finish timestamps in current artifacts.",
            "needed_next": "Add request dispatch and finish/error timestamps to future replayer/proxy metrics.",
            "reviewer_risk": "Cannot use these runs to prove online per-session sequentiality.",
        },
        {
            "claim": "Batch 1 workload shape can be characterized from formatted traces and metrics.",
            "status": "supported_for_trace_shape",
            "supporting_data": "analysis/characterization/analyze.py outputs workload_summary/session_skew/KV footprint when given trace and kv_bytes_per_token.",
            "needed_next": "Run on dash0 compact formatted trace for canonical full-trace numbers.",
            "reviewer_risk": "Actual cache reuse decomposition needs cached_tokens joined with hash_ids.",
        },
        {
            "claim": "Static PD separation is worse than combined in existing 200-request GPU A/B.",
            "status": "supported_by_existing_artifact" if has_metrics else "unavailable",
            "supporting_data": "outputs/gpu_ab_combined vs outputs/gpu_ab_pdsep metrics.summary.json.",
            "needed_next": "Refresh with PD matrix, multiple seeds, cudagraph-enabled methodology.",
            "reviewer_risk": "Legacy run has no per-stage TTFT breakdown and no step-level KV occupancy.",
        },
        {
            "claim": "Elastic transfer-based migration does not improve high-contention 500-request run.",
            "status": "supported_by_existing_artifact" if has_metrics else "unavailable",
            "supporting_data": "outputs/contention_16s_ts10 vs outputs/contention_16s_elastic metrics.summary.json and gpu_util.csv.",
            "needed_next": "Attribute whether failure is trigger quality, transfer overhead, or wrong load regime.",
            "reviewer_risk": "Existing metrics lack actual sequentiality proof and per-request transfer waterfall.",
        },
        {
            "claim": "PD-colo prefill/decode interference is not yet directly proven by step-level data in this package.",
            "status": "not_yet_supported",
            "supporting_data": "No decode-step and prefill-overlap timestamp artifact found in summarized runs.",
            "needed_next": "Run Batch 2 controlled same-worker/different-worker injection with step timestamps.",
            "reviewer_risk": "Cannot claim interference as causal without Batch 2.",
        },
        {
            "claim": "Session hot-spot residual imbalance is suggested but not fully attributed.",
            "status": "partially_supported" if has_gpu else "unavailable",
            "supporting_data": "gpu_util.csv shows per-GPU mean-util imbalance in existing runs.",
            "needed_next": "Collect per-worker queue delay, session-to-worker map, and per-session token mass per worker.",
            "reviewer_risk": "GPU util imbalance alone is not enough to prove session hot-spot.",
        },
        {
            "claim": "SRR is not measured by existing fixed-request runs.",
            "status": "not_yet_supported",
            "supporting_data": "No arrival-rate sweep artifacts found.",
            "needed_next": "Implement Batch 4 Poisson session-arrival SRR sweep.",
            "reviewer_risk": "Latency-at-one-load cannot support sustainable throughput claim.",
        },
    ]


def build_risk_register(summaries: list[JsonDict]) -> list[JsonDict]:
    return [
        {
            "risk": "Session sequentiality not proven",
            "severity": "high",
            "evidence": "Current metrics include trace timestamp and latency but not actual dispatch/finish wall-clock timestamps.",
            "mitigation": "Add dispatch/finish timestamps and run Batch 0 before SRR claims.",
        },
        {
            "risk": "Legacy PD-sep data may not match final methodology",
            "severity": "medium",
            "evidence": "PD matrix scaffold exists separately; some old runs used earlier flags/methodology.",
            "mitigation": "Use fresh PD matrix for paper-grade claims.",
        },
        {
            "risk": "GPU util is not a sufficient hot-spot proof",
            "severity": "medium",
            "evidence": "Existing artifacts have gpu_util.csv but lack per-worker queue and session ownership.",
            "mitigation": "Add route-decision and per-worker queue logs for Batch 3.",
        },
        {
            "risk": "Cache reuse decomposition is incomplete without joined hash/cache-hit data",
            "severity": "medium",
            "evidence": "Trace has hash_ids; metrics have cached_tokens; request IDs may not join across all artifacts.",
            "mitigation": "Emit hash_ids/session_id/cached_tokens in the same per-request record.",
        },
    ]


def render_current_results(
    summaries: list[JsonDict],
    comparisons: list[JsonDict],
    claim_matrix: list[JsonDict],
    risk_register: list[JsonDict],
) -> str:
    lines = [
        "# Current Characterization Results",
        "",
        f"Generated: {dt.datetime.now(dt.timezone.utc).isoformat()}",
        f"Git commit: `{git_commit()}`",
        "",
        "## Existing Run Summaries",
        "",
        "| Run | OK/Req | TTFT p50/p90 | E2E p50/p90 | TPOT p90 | GPU mean util | GPU imbalance |",
        "|---|---:|---:|---:|---:|---:|---:|",
    ]
    for s in summaries:
        lines.append(
            "| {run} | {ok}/{req} | {ttft50}/{ttft90} | {e2e50}/{e2e90} | {tpot90} | {gpu_mean} | {gpu_imb} |".format(
                run=s["run"],
                ok=fmt(s.get("success_count")),
                req=fmt(s.get("request_count")),
                ttft50=fmt(stat_value(s, "ttft_stats_s", "p50")),
                ttft90=fmt(stat_value(s, "ttft_stats_s", "p90")),
                e2e50=fmt(stat_value(s, "latency_stats_s", "p50")),
                e2e90=fmt(stat_value(s, "latency_stats_s", "p90")),
                tpot90=fmt(stat_value(s, "tpot_stats_s", "p90")),
                gpu_mean=fmt(nested(s, ["gpu_summary", "mean_util_pct"])),
                gpu_imb=fmt(nested(s, ["gpu_summary", "max_min_ratio"])),
            )
        )
    lines.extend([
        "",
        "## Pairwise Comparisons",
        "",
        "| Comparison | TTFT p50 Δ | TTFT p90 Δ | E2E p50 Δ | E2E p90 Δ | TPOT p90 Δ | Wall-clock Δ |",
        "|---|---:|---:|---:|---:|---:|---:|",
    ])
    for c in comparisons:
        lines.append(
            "| {name} | {ttft50} | {ttft90} | {e2e50} | {e2e90} | {tpot90} | {wall} |".format(
                name=c["name"],
                ttft50=fmt_pct(c.get("ttft_p50_delta_pct")),
                ttft90=fmt_pct(c.get("ttft_p90_delta_pct")),
                e2e50=fmt_pct(c.get("e2e_p50_delta_pct")),
                e2e90=fmt_pct(c.get("e2e_p90_delta_pct")),
                tpot90=fmt_pct(c.get("tpot_p90_delta_pct")),
                wall=fmt_pct(c.get("wall_clock_delta_pct")),
            )
        )
    lines.extend([
        "",
        "## What We Can Say Now",
        "",
    ])
    for item in claim_matrix:
        lines.append(f"- **{item['status']}**: {item['claim']}")
        lines.append(f"  Supporting data: {item['supporting_data']}")
        lines.append(f"  Next: {item['needed_next']}")
    lines.extend([
        "",
        "## Main Reviewer Risks",
        "",
    ])
    for item in risk_register:
        lines.append(f"- **{item['severity']}**: {item['risk']} - {item['mitigation']}")
    lines.append("")
    return "\n".join(lines)


def render_claim_matrix(items: list[JsonDict]) -> str:
    lines = [
        "# Characterization Claim Matrix",
        "",
        "| Claim | Status | Supporting Data | Needed Next | Reviewer Risk |",
        "|---|---|---|---|---|",
    ]
    for item in items:
        lines.append(
            f"| {item['claim']} | `{item['status']}` | {item['supporting_data']} | {item['needed_next']} | {item['reviewer_risk']} |"
        )
    lines.append("")
    return "\n".join(lines)


def render_risk_register(items: list[JsonDict]) -> str:
    lines = [
        "# Reviewer Risk Register",
        "",
        "| Risk | Severity | Evidence | Mitigation |",
        "|---|---|---|---|",
    ]
    for item in items:
        lines.append(
            f"| {item['risk']} | `{item['severity']}` | {item['evidence']} | {item['mitigation']} |"
        )
    lines.append("")
    return "\n".join(lines)


def render_figures_index(summaries: list[JsonDict]) -> str:
    return "\n".join([
        "# Figures Index",
        "",
        "No generated figures are committed by this script. Batch-specific figures should be generated from:",
        "",
        "- `analysis/characterization/analyze.py` for Batch 0/1 trace figures.",
        "- future Batch 2 step-timeline artifacts for interference plots.",
        "- future Batch 3 per-worker/session artifacts for hot-spot plots.",
        "- future Batch 4 arrival-rate sweep artifacts for SRR curves.",
        "",
        "This file exists so the audit package has a stable placeholder until fresh figures are generated.",
        "",
    ])


def render_reproduction_commands(args: argparse.Namespace, run_dirs: list[Path]) -> str:
    runs = " ".join(str(p) for p in run_dirs)
    return "\n".join([
        "#!/usr/bin/env bash",
        "set -euo pipefail",
        "",
        "# Rebuild this current-results audit package.",
        f"python3 analysis/characterization/summarize_runs.py --output-dir {args.output_dir} --runs {runs}",
        "",
        "# Example Batch 0/1 local trace analysis.",
        "python3 analysis/characterization/analyze.py \\",
        "  --trace traces/w600_r0.0015_st30.jsonl \\",
        "  --kv-bytes-per-token 98304 \\",
        "  --task-name w600_local_full_trace \\",
        "  --overwrite",
        "",
    ])


def load_json(path: Path) -> JsonDict:
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}
    return data if isinstance(data, dict) else {}


def load_jsonl(path: Path) -> list[JsonDict]:
    if not path.exists():
        return []
    rows = []
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            if not line.strip():
                continue
            try:
                row = json.loads(line)
            except Exception:
                continue
            if isinstance(row, dict):
                rows.append(row)
    return rows


def write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def first_present(row: JsonDict, keys: list[str]) -> Any:
    for key in keys:
        if key in row:
            return row[key]
    return None


def stat_value(run: JsonDict, stat_key: str, value_key: str) -> float | None:
    stats_obj = run.get(stat_key)
    if not isinstance(stats_obj, dict):
        return None
    return to_float(stats_obj.get(value_key))


def nested(row: JsonDict, keys: list[str]) -> Any:
    cur: Any = row
    for key in keys:
        if not isinstance(cur, dict):
            return None
        cur = cur.get(key)
    return cur


def pct_delta(base: Any, variant: Any) -> float | None:
    b = to_float(base)
    v = to_float(variant)
    if b is None or v is None or b == 0:
        return None
    return (v - b) / b * 100.0


def to_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        out = float(value)
    except (TypeError, ValueError):
        return None
    return out if math.isfinite(out) else None


def stats(values: list[float]) -> JsonDict | None:
    clean = sorted(float(v) for v in values if math.isfinite(float(v)))
    if not clean:
        return None
    return {
        "count": len(clean),
        "mean": statistics.fmean(clean),
        "p50": percentile(clean, 0.50),
        "p90": percentile(clean, 0.90),
        "p95": percentile(clean, 0.95),
        "p99": percentile(clean, 0.99),
        "max": clean[-1],
    }


def percentile(values: list[float], q: float) -> float:
    if len(values) == 1:
        return values[0]
    rank = q * (len(values) - 1)
    lo = int(rank)
    hi = min(lo + 1, len(values) - 1)
    frac = rank - lo
    return values[lo] * (1 - frac) + values[hi] * frac


def top_contribution(values: list[float]) -> JsonDict:
    clean = sorted([v for v in values if math.isfinite(v)], reverse=True)
    total = sum(clean)
    if not clean or total <= 0:
        return {"top_1pct": None, "top_5pct": None, "top_10pct": None}

    def frac(pct: float) -> float:
        k = max(1, math.ceil(len(clean) * pct))
        return sum(clean[:k]) / total

    return {
        "top_1pct": frac(0.01),
        "top_5pct": frac(0.05),
        "top_10pct": frac(0.10),
    }


def fmt(value: Any) -> str:
    num = to_float(value)
    if num is None:
        return "n/a"
    if abs(num - round(num)) < 1e-9 and abs(num) < 1_000_000:
        return str(int(round(num)))
    return f"{num:.3g}"


def fmt_pct(value: Any) -> str:
    num = to_float(value)
    if num is None:
        return "n/a"
    return f"{num:+.1f}%"


def git_commit() -> str:
    try:
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
        )
    except Exception:
        return ""
    return result.stdout.strip()


if __name__ == "__main__":
    main()
