"""How long Loupe takes: render and judge time per page on the benchmark's untouched pages, on this machine's CPU.

Run it when nothing else is using the machine. Render time is a fresh page load, so it includes the network.

Usage: timing.py [number of pages, default 12]
"""

import json
import statistics
import subprocess
import sys
import tempfile
import time
from pathlib import Path

sys.stdout.reconfigure(encoding="utf-8", errors="replace")
PACKAGE = Path(__file__).resolve().parent
sys.path.insert(0, str(PACKAGE))

import benchmark
import page_audit
from safetensors.torch import load_file


def main():
    count = int(sys.argv[1]) if len(sys.argv) > 1 else 12
    read = page_audit.load_read()
    brief = dict(page_audit.DEFAULT_BRIEF)
    rows = []
    for name, url in list(benchmark.PAGES.items())[::max(1, len(benchmark.PAGES) // count)][:count]:
        directory = Path(tempfile.mkdtemp(prefix="loupe-timing-"))
        started = time.perf_counter()
        completed = subprocess.run(["node", str(PACKAGE / "collect_page.cjs"), url, str(directory), "1440", "plain"], capture_output=True, timeout=240)
        rendered = time.perf_counter()
        if completed.returncode != 0:
            continue
        results = page_audit.judge(page_audit.measure(directory, "dom"), brief, read)
        judged = time.perf_counter()
        rows.append({"page": name, "render_seconds": rendered - started, "judge_seconds": judged - rendered, "checks": len(results)})
        print(name, f"render {rendered - started:.2f} s, judge {judged - rendered:.2f} s, {len(results)} checks", flush=True)
    weights = PACKAGE / "read_v2" / "read.safetensors"
    summary = {"pages": len(rows), "render_seconds_median": statistics.median(row["render_seconds"] for row in rows),
               "judge_seconds_median": statistics.median(row["judge_seconds"] for row in rows),
               "judge_seconds_max": max(row["judge_seconds"] for row in rows), "checks_median": statistics.median(row["checks"] for row in rows),
               "parameters": int(sum(tensor.numel() for tensor in load_file(str(weights)).values())), "weights_kib": weights.stat().st_size / 1024,
               "device": "CPU", "rows": rows}
    (benchmark.OUTPUT / "timing.json").write_text(json.dumps(summary, indent=1), encoding="utf-8")
    print({key: value for key, value in summary.items() if key != "rows"})


if __name__ == "__main__":
    main()