"""Parse surya_timing_summary lines from logs/info.log and aggregate spans.

Usage:
    python scripts/parse_timing.py logs/info.log
"""
from __future__ import annotations

import argparse
import ast
import json
import re
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional

_BATCH_RE = re.compile(r"batch_size=(\d+)")
_REQID_RE = re.compile(r"request_id=(\S+)")


def parse_line(line: str) -> Optional[Dict[str, Any]]:
    if "surya_timing_summary" not in line or "events=" not in line:
        return None
    events_str = line.split("events=", 1)[1].strip()
    try:
        events = ast.literal_eval(events_str)
    except (ValueError, SyntaxError):
        return None
    batch_match = _BATCH_RE.search(line)
    reqid_match = _REQID_RE.search(line)
    return {
        "request_id": reqid_match.group(1) if reqid_match else "-",
        "batch_size": int(batch_match.group(1)) if batch_match else 0,
        "events": events,
    }


def aggregate(records: List[Dict[str, Any]]) -> Dict[str, Dict[str, float]]:
    stats: Dict[str, Dict[str, float]] = defaultdict(
        lambda: {"count": 0, "total_ms": 0.0, "mean_ms": 0.0, "total_tokens": 0}
    )
    for rec in records:
        for event in rec["events"]:
            s = stats[event["name"]]
            s["count"] += 1
            s["total_ms"] += event.get("duration_ms", 0.0)
            tokens = (event.get("metadata") or {}).get("token_count")
            if tokens:
                s["total_tokens"] += tokens
    for s in stats.values():
        s["mean_ms"] = round(s["total_ms"] / s["count"], 2) if s["count"] else 0.0
        s["total_ms"] = round(s["total_ms"], 2)
    return dict(stats)


def main() -> None:
    parser = argparse.ArgumentParser(description="Aggregate surya timing spans from a log file.")
    parser.add_argument("logfile", type=Path)
    args = parser.parse_args()
    records = [
        rec
        for rec in (parse_line(line) for line in args.logfile.read_text(encoding="utf-8").splitlines())
        if rec is not None
    ]
    print(json.dumps({"records": len(records), "spans": aggregate(records)}, indent=2))


if __name__ == "__main__":
    main()