#!/usr/bin/env python3
"""Compare the app's debug diagnostics with the budgets in tools/budgets.json.

Debug builds print lines like
    DIAG {"event":"first_frame","startup_ms":123,"mem_used":4567,"mem_total":786432}
(source/util/Diag.mc). app-diag.sh runs the app on each test device, saves the
console output to build/diag/<device>.log, then calls this script.

Usage:  check_budgets.py {perf,mem} LOG [LOG ...] [--budgets FILE] [--summary OUT.json]
Exit:   0 within budget, 1 over budget, 2 missing or unreadable diagnostics.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

EXIT_OK, EXIT_OVER, EXIT_MISSING = 0, 1, 2


def read_events(text: str) -> list[dict[str, Any]]:
    """Every DIAG event in a console log, in order. Malformed lines are skipped."""
    events = []
    for line in text.splitlines():
        start = line.find("DIAG {")
        if start < 0:
            continue
        try:
            event = json.loads(line[start + len("DIAG ") :])
        except json.JSONDecodeError:
            continue
        if isinstance(event, dict):
            events.append(event)
    return events


def measure(kind: str, events: list[dict[str, Any]]) -> dict[str, float]:
    """The values a budget kind looks at, from the first first_frame event."""
    frame = next((e for e in events if e.get("event") == "first_frame"), None)
    if frame is None:
        return {}
    if kind == "perf":
        return {"startup_ms": float(frame["startup_ms"])}
    total = float(frame["mem_total"])
    return {"mem_used_fraction": float(frame["mem_used"]) / total if total else 1.0}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Check diagnostics against budgets.")
    parser.add_argument("kind", choices=("perf", "mem"))
    parser.add_argument("logs", type=Path, nargs="+")
    parser.add_argument("--budgets", type=Path, default=Path("tools/budgets.json"))
    parser.add_argument("--summary", type=Path)
    args = parser.parse_args(argv)

    budgets: dict[str, float] = json.loads(args.budgets.read_text(encoding="utf-8"))[args.kind]
    worst = EXIT_OK
    summary: dict[str, dict[str, float]] = {}
    for log in args.logs:
        device = log.stem
        try:
            values = measure(args.kind, read_events(log.read_text(encoding="utf-8", errors="replace")))
        except (OSError, KeyError, ValueError) as err:
            print(f"[{device}] cannot read diagnostics: {err}")
            worst = max(worst, EXIT_MISSING)
            continue
        if not values:
            print(f"[{device}] no DIAG first_frame line (did the app start? is this a debug build?)")
            worst = max(worst, EXIT_MISSING)
            continue
        summary[device] = values
        for name, limit in budgets.items():
            value = values[name]
            over = value > limit
            print(f"[{device}] {name} = {value:.3g} (budget {limit:.3g}) {'OVER BUDGET' if over else 'ok'}")
            if over:
                worst = max(worst, EXIT_OVER)
    if args.summary:
        args.summary.parent.mkdir(parents=True, exist_ok=True)
        args.summary.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return worst


if __name__ == "__main__":
    sys.exit(main())