#!/usr/bin/env python3
"""Summarize workload-regime Frontier metrics and closed-loop batch state."""

from __future__ import annotations

import argparse
import json
import math
from collections import defaultdict
from pathlib import Path
from typing import Any


METRICS = (
    "ttft_mean_ms",
    "ttft_p90_ms",
    "tpot_mean_ms",
    "tpot_p90_ms",
    "e2e_mean_ms",
    "e2e_p90_ms",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--sim-root", type=Path, required=True)
    parser.add_argument("--workload-manifest", type=Path, required=True)
    parser.add_argument("--json-output", type=Path, required=True)
    parser.add_argument("--markdown-output", type=Path, required=True)
    return parser.parse_args()


def summarize_ledger(path: Path) -> dict[str, float | int | None]:
    total_time = 0.0
    weighted_batch = 0.0
    decode_time = 0.0
    decode_weighted_batch = 0.0
    prefill_time = 0.0
    prefill_weighted_batch = 0.0
    max_batch = 0
    batches = 0
    moe_time = 0.0
    collective_time = 0.0
    for line in path.open():
        if not line.strip():
            continue
        row = json.loads(line)
        request_tokens = [int(value) for value in row["request_num_tokens"]]
        batch = len(request_tokens)
        execution = row["execution_time"]
        duration = float(execution["total_time_ms"])
        components = execution["component_ledger_ms"]
        total_time += duration
        weighted_batch += batch * duration
        max_batch = max(max_batch, batch)
        batches += 1
        if request_tokens and all(value == 1 for value in request_tokens):
            decode_time += duration
            decode_weighted_batch += batch * duration
        else:
            prefill_time += duration
            prefill_weighted_batch += batch * duration
        moe_time += sum(
            float(components.get(name, 0.0))
            for name in (
                "moe_gating_linear_time",
                "moe_gating_routing_topk_time",
                "moe_grouped_gemm_time",
                "moe_shuffling_time",
            )
        )
        collective_time += sum(
            float(components.get(name, 0.0))
            for name in (
                "attention_all_reduce_time",
                "mlp_all_reduce_time",
                "moe_tensor_parallel_allgather_time",
                "expert_parallel_communication_time",
            )
        )
    return {
        "batches": batches,
        "max_batch_size": max_batch,
        "time_weighted_batch_size": weighted_batch / total_time if total_time else None,
        "decode_time_weighted_batch_size": (
            decode_weighted_batch / decode_time if decode_time else None
        ),
        "prefill_mixed_time_weighted_batch_size": (
            prefill_weighted_batch / prefill_time if prefill_time else None
        ),
        "decode_time_fraction": decode_time / total_time if total_time else None,
        "moe_time_fraction": moe_time / total_time if total_time else None,
        "collective_time_fraction": (
            collective_time / total_time if total_time else None
        ),
    }


def find_one(root: Path, pattern: str) -> Path:
    matches = list(root.glob(pattern))
    if len(matches) != 1:
        raise ValueError(f"expected one {pattern} under {root}, got {matches}")
    return matches[0]


def load_rows(sim_root: Path, workload_manifest: Path) -> list[dict[str, Any]]:
    suite = json.loads(workload_manifest.read_text())
    cases = {case["public_csv_sha256"]: case for case in suite["cases"]}
    rows = []
    for result_path in sorted(sim_root.glob("tp*-prefix-*/runs/*/*/result.json")):
        result = json.loads(result_path.read_text())
        case = cases.get(result.get("trace_sha256"))
        if case is None:
            raise ValueError(f"unknown trace hash in {result_path}")
        row: dict[str, Any] = {
            "family": case["family"],
            "rho": case["rho"],
            "prefix_caching": case["prefix_caching"],
            "tp": result["config"]["tp"],
            "mns": result["config"]["mns"],
            "config": result["config"]["name"],
            "status": result["status"],
            "failure_kind": result.get("failure_kind"),
            "result_path": str(result_path),
        }
        if result["status"] == "completed":
            row.update({name: result["score"].get(name) for name in METRICS})
            ledger = find_one(
                result_path.parent,
                "metrics/**/frontier_stage_batch_ledger.jsonl",
            )
            row["ledger_path"] = str(ledger)
            row.update(summarize_ledger(ledger))
        rows.append(row)
    return rows


def best_configs(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    groups: dict[tuple[str, float], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        if row["status"] == "completed":
            groups[(row["family"], row["rho"])].append(row)
    best = []
    for (family, rho), records in sorted(groups.items()):
        output: dict[str, Any] = {"family": family, "rho": rho}
        for metric in METRICS:
            valid = [row for row in records if row.get(metric) is not None]
            if valid:
                winner = min(valid, key=lambda row: (row[metric], row["config"]))
                output[metric] = {
                    "config": winner["config"],
                    "value_ms": winner[metric],
                }
        best.append(output)
    return best


def state_knees(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        value = row.get("decode_time_weighted_batch_size")
        if row["status"] == "completed" and value is not None:
            groups[(row["family"], row["config"])].append(row)
    knees = []
    for (family, config), records in sorted(groups.items()):
        records.sort(key=lambda row: row["rho"])
        candidates = []
        for left, right in zip(records, records[1:]):
            delta = (
                right["decode_time_weighted_batch_size"]
                - left["decode_time_weighted_batch_size"]
            )
            candidates.append((abs(delta), delta, left, right))
        if candidates:
            _, delta, left, right = max(candidates, key=lambda item: item[0])
            knees.append(
                {
                    "family": family,
                    "config": config,
                    "rho_interval": [left["rho"], right["rho"]],
                    "decode_batch_before": left[
                        "decode_time_weighted_batch_size"
                    ],
                    "decode_batch_after": right[
                        "decode_time_weighted_batch_size"
                    ],
                    "decode_batch_delta": delta,
                }
            )
    return knees


def render_markdown(summary: dict[str, Any]) -> str:
    lines = [
        "# Frontier simulator workload-regime discovery",
        "",
        f"- Expected cells: {summary['coverage']['expected']}",
        f"- Observed cells: {summary['coverage']['observed']}",
        f"- Completed: {summary['coverage']['completed']}",
        f"- Failed: {summary['coverage']['failed']}",
        "",
        "This table reports simulator predictions only. It does not classify a workload as reliable without matched real-system evidence.",
        "",
        "## Simulator winners",
        "",
        "| Family | rho | TTFT mean | TTFT p90 | TPOT mean | TPOT p90 |",
        "|---|---:|---|---|---|---|",
    ]
    for row in summary["best_configs"]:
        values = []
        for metric in (
            "ttft_mean_ms",
            "ttft_p90_ms",
            "tpot_mean_ms",
            "tpot_p90_ms",
        ):
            winner = row.get(metric)
            values.append(
                "—"
                if winner is None
                else f"{winner['config']} ({winner['value_ms']:.2f} ms)"
            )
        lines.append(
            f"| {row['family']} | {row['rho']:.2f} | " + " | ".join(values) + " |"
        )
    lines.extend(
        [
            "",
            "## Largest decode-batch transition per config",
            "",
            "| Family | Config | rho interval | Batch before | Batch after | Delta |",
            "|---|---|---|---:|---:|---:|",
        ]
    )
    for row in summary["state_knees"]:
        lines.append(
            "| {family} | {config} | {left:.2f}–{right:.2f} | {before:.2f} | {after:.2f} | {delta:+.2f} |".format(
                family=row["family"],
                config=row["config"],
                left=row["rho_interval"][0],
                right=row["rho_interval"][1],
                before=row["decode_batch_before"],
                after=row["decode_batch_after"],
                delta=row["decode_batch_delta"],
            )
        )
    return "\n".join(lines) + "\n"


def main() -> None:
    args = parse_args()
    rows = load_rows(args.sim_root, args.workload_manifest)
    completed = sum(row["status"] == "completed" for row in rows)
    summary = {
        "schema": "frontier-workload-regime-simulator-analysis-v1",
        "coverage": {
            "expected": 7 * 5 * 12,
            "observed": len(rows),
            "completed": completed,
            "failed": len(rows) - completed,
        },
        "rows": rows,
        "best_configs": best_configs(rows),
        "state_knees": state_knees(rows),
    }
    for output in (args.json_output, args.markdown_output):
        output.parent.mkdir(parents=True, exist_ok=True)
    args.json_output.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
    args.markdown_output.write_text(render_markdown(summary))
    print(args.markdown_output)


if __name__ == "__main__":
    main()