#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

DEFAULT_RIRS = ["afrinic", "apnic", "arin", "lacnic", "ripe"]


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Aggregate per-RIR replay case reports")
    p.add_argument("--case-root", required=True, help="directory containing <rir>/<rir>_case_report.json")
    p.add_argument("--out-md", required=True)
    p.add_argument("--out-json", required=True)
    p.add_argument("--rirs", nargs="*", default=None, help="RIRs to include (default: all 5)")
    return p.parse_args()


def read_case(case_root: Path, rir: str) -> dict:
    path = case_root / rir / f"{rir}_case_report.json"
    return json.loads(path.read_text(encoding="utf-8"))


def geomean(values: list[float]) -> float:
    vals = [v for v in values if v > 0]
    if not vals:
        return 0.0
    return math.exp(sum(math.log(v) for v in vals) / len(vals))


def build_summary(cases: list[dict]) -> dict:
    snapshot_ratios = [c["snapshot"]["ratio"] for c in cases]
    delta_ratios = [c["delta"]["ratio"] for c in cases]
    return {
        "cases": cases,
        "summary": {
            "snapshot_all_match": all(c["snapshot"]["match"] for c in cases),
            "delta_all_match": all(c["delta"]["match"] for c in cases),
            "snapshot_ratio_geomean": geomean(snapshot_ratios),
            "delta_ratio_geomean": geomean(delta_ratios),
            "all_ratio_geomean": geomean(snapshot_ratios + delta_ratios),
        },
    }


def write_md(path: Path, data: dict) -> None:
    lines = []
    lines.append("# Multi-RIR Replay Summary\n\n")
    lines.append("## Correctness + Timing\n\n")
    lines.append("| RIR | snapshot_match | snapshot_ours_s | snapshot_routinator_s | snapshot_ratio | delta_match | delta_ours_s | delta_routinator_s | delta_ratio |\n")
    lines.append("|---|---|---:|---:|---:|---|---:|---:|---:|\n")
    for case in data["cases"]:
        lines.append(
            f"| {case['rir']} | {str(case['snapshot']['match']).lower()} | {case['snapshot']['ours_seconds']:.3f} | {case['snapshot']['routinator_seconds']:.3f} | {case['snapshot']['ratio']:.3f} | {str(case['delta']['match']).lower()} | {case['delta']['ours_seconds']:.3f} | {case['delta']['routinator_seconds']:.3f} | {case['delta']['ratio']:.3f} |\n"
        )
    s = data["summary"]
    lines.append("\n## Aggregate Metrics\n\n")
    lines.append("| metric | value |\n")
    lines.append("|---|---:|\n")
    lines.append(f"| snapshot_all_match | {str(s['snapshot_all_match']).lower()} |\n")
    lines.append(f"| delta_all_match | {str(s['delta_all_match']).lower()} |\n")
    lines.append(f"| snapshot_ratio_geomean | {s['snapshot_ratio_geomean']:.3f} |\n")
    lines.append(f"| delta_ratio_geomean | {s['delta_ratio_geomean']:.3f} |\n")
    lines.append(f"| all_ratio_geomean | {s['all_ratio_geomean']:.3f} |\n")
    path.write_text("".join(lines), encoding="utf-8")


def main() -> int:
    args = parse_args()
    case_root = Path(args.case_root)
    rirs = args.rirs or DEFAULT_RIRS
    cases = [read_case(case_root, rir) for rir in rirs]
    data = build_summary(cases)
    out_md = Path(args.out_md)
    out_json = Path(args.out_json)
    out_md.parent.mkdir(parents=True, exist_ok=True)
    out_json.parent.mkdir(parents=True, exist_ok=True)
    out_json.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    write_md(out_md, data)
    print(out_md)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
