from __future__ import annotations

import json
import statistics
from pathlib import Path


ROOT = Path(__file__).resolve().parent
RESULTS = ROOT / ".work" / "results"
MIB = 1024 * 1024


def is_worker_staged(result: dict) -> bool:
    return (
        # Model-execution A/B runs use their own phase-level report. Keep this
        # runtime-only summary on the original fixed cohort without a decoder
        # query so later candidate measurements cannot silently change it.
        result.get("decoderMode") is None
        and result.get("benchmarkScope") in (None, "full")
        and any(
            sample.get("label") == "vision-worker-session-ready"
            for sample in result.get("memorySamples", [])
        )
    )


def is_resident_control(result: dict) -> bool:
    return (
        result.get("memoryMode") == "default"
        and result.get("loaderMode") == "bytes"
        and not is_worker_staged(result)
        and result.get("scheduleMode", "resident") == "resident"
    )


def median(rows: list[dict], value) -> float:
    return statistics.median(value(row) for row in rows)


def main() -> None:
    loaded = [json.loads(path.read_text(encoding="utf-8")) for path in RESULTS.glob("*.json")]
    for tier in ("121", "242"):
        controls = [row for row in loaded if row.get("tier") == tier and is_resident_control(row)]
        staged = [row for row in loaded if row.get("tier") == tier and is_worker_staged(row)]
        if not controls or not staged:
            print(f"tier {tier}: missing control or worker-staged results")
            continue
        control_outputs = [item["actual"] for item in controls[0]["results"]]
        if any([item["actual"] for item in row["results"]] != control_outputs for row in staged):
            raise RuntimeError(f"tier {tier}: optimized OCR output differs from control")

        control_renderer = median(
            controls, lambda row: row["processMemory"]["renderer"]["peakDeltaBytes"] / MIB
        )
        staged_renderer = median(
            staged, lambda row: row["processMemory"]["renderer"]["peakDeltaBytes"] / MIB
        )
        control_heap = median(
            controls,
            lambda row: max(sample.get("usedJSHeapSize", 0) for sample in row["memorySamples"]) / MIB,
        )
        staged_heap = median(
            staged,
            lambda row: max(sample.get("usedJSHeapSize", 0) for sample in row["memorySamples"]) / MIB,
        )
        print(
            f"tier {tier}: parity=PASS controls={len(controls)} staged={len(staged)} "
            f"renderer={control_renderer:.1f}->{staged_renderer:.1f} MiB "
            f"({(staged_renderer / control_renderer - 1) * 100:+.1f}%), "
            f"js_heap={control_heap:.1f}->{staged_heap:.1f} MiB "
            f"({(staged_heap / control_heap - 1) * 100:+.1f}%)"
        )


if __name__ == "__main__":
    main()