"""Establish a repeated canonical baseline through the resident H3 API."""

from __future__ import annotations

import argparse
import hashlib
import json
import statistics
import time
from pathlib import Path
from urllib.request import Request, urlopen


DEFAULT_PROMPT = (
    "A playful orange tabby cat starts in an ordinary cozy living room in a normal house, "
    "afternoon light, sofa and rug. The cat crouches, jumps, and does one clean athletic "
    "backflip in slow motion. As the backflip completes there is a sharp cinematic cut: "
    "the cat lands perfectly on a glowing neon disco dance floor wearing oversized black "
    "sunglasses. Mirror ball reflections, colorful lights, joyful party energy, stylish "
    "and funny, clear before-and-after transformation."
)
EXPECTED_VIDEO_SHA256 = "c62d23a42972eab907ba42f93c50247ff17a9c454b4a53fe93d2e34f9fefe578"
EXPECTED_AUDIO_SHA256 = "852005383770480a6503504e1ffec86dd1fb63a69c6400f92da18e39e0986de2"


def get_json(url: str, timeout: float = 30.0) -> dict:
    with urlopen(url, timeout=timeout) as response:
        return json.loads(response.read().decode())


def post_json(url: str, payload: dict, timeout: float) -> dict:
    request = Request(
        url,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode())


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--server", default="http://127.0.0.1:8001")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--timeout", type=float, default=1200.0)
    parser.add_argument("--source-commit", default="a29b8960b0f887c20e74dafa16a24c37d6508b4e")
    parser.add_argument("--image", required=True)
    args = parser.parse_args()
    if args.runs < 3:
        raise ValueError("the authoritative baseline requires at least three measured runs")

    ready_before = get_json(f"{args.server}/ready")
    if not ready_before.get("ready"):
        raise RuntimeError("resident runtime is not ready")
    if ready_before["runtime"]["current_attention"] != "sage2":
        raise RuntimeError("resident runtime is not using Sage2")
    if not ready_before["runtime"]["fc2_lt"]["enabled"]:
        raise RuntimeError("guarded FC2 schedule is disabled")

    def payload(label: str) -> dict:
        return {
            "prompt": DEFAULT_PROMPT,
            "output": f"/output/h3-blackwell-runtime/post-fc2-{label}.mp4",
            "width": 1344,
            "height": 768,
            "frames": 124,
            "steps": 12,
            "seed": 440420,
            "attention": "sage2",
            "turbo": None,
            "cache_mode": None,
            "mux_audio": True,
            "keep_intermediates": False,
            "benchmark_text_tokens": 100,
        }

    warmup = post_json(f"{args.server}/generate", payload("canonical-warmup"), args.timeout)
    measured = []
    for index in range(1, args.runs + 1):
        started = time.perf_counter()
        response = post_json(
            f"{args.server}/generate", payload(f"canonical-run-{index}"), args.timeout,
        )
        response["client_wall_seconds"] = time.perf_counter() - started
        measured.append(response)

    for name, response in [("warmup", warmup), *[
        (f"run_{index}", value) for index, value in enumerate(measured, start=1)
    ]]:
        dispatch = response["fc2_dispatch_delta"]
        if dispatch != {"attempts": 600, "successes": 600, "fallbacks": 0}:
            raise RuntimeError(f"{name} FC2 dispatch validation failed: {dispatch}")
        if len(response["sampling_steps"]) != 12:
            raise RuntimeError(f"{name} did not report 12 sampling steps")

    checksum_pairs = [
        (row["latent_checksums"]["video_sha256"], row["latent_checksums"]["audio_sha256"])
        for row in measured
    ]
    exact_parity = len(set(checksum_pairs)) == 1
    if not exact_parity:
        raise RuntimeError(f"measured latent checksums differ: {checksum_pairs}")
    expected_checksums = (EXPECTED_VIDEO_SHA256, EXPECTED_AUDIO_SHA256)
    if checksum_pairs[0] != expected_checksums:
        raise RuntimeError(
            f"latent checksums differ from canonical reference: "
            f"expected {expected_checksums}, got {checksum_pairs[0]}"
        )

    sampling_seconds = [row["sampling_seconds"] for row in measured]
    report = {
        "name": "gb10-post-fc2-resident-baseline",
        "source_commit": args.source_commit,
        "image": args.image,
        "measurement_date": "2026-08-26",
        "workload": {
            "resolution": [1344, 768],
            "frames": 124,
            "steps": 12,
            "seed": 440420,
            "attention": "sage2",
            "prompt": DEFAULT_PROMPT,
            "prompt_sha256": hashlib.sha256(DEFAULT_PROMPT.encode()).hexdigest(),
        },
        "measurement_policy": {
            "canonical_warmup_runs": 1,
            "measured_runs": args.runs,
            "authoritative_timing": "median synchronized resident sampling_seconds",
            "step_timing": "deferred CUDA event elapsed time; no per-step synchronization",
            "profiling": False,
        },
        "sampling_seconds": sampling_seconds,
        "median_sampling_seconds": statistics.median(sampling_seconds),
        "sampling_steps": [row["sampling_steps"] for row in measured],
        "sampling_peak_allocated_bytes": [
            row["sampling_peak_allocated_bytes"] for row in measured
        ],
        "sampling_peak_reserved_bytes": [
            row["sampling_peak_reserved_bytes"] for row in measured
        ],
        "latent_checksums": {
            "video_sha256": checksum_pairs[0][0],
            "audio_sha256": checksum_pairs[0][1],
            "exact_across_measured_runs": exact_parity,
            "matches_canonical_reference": True,
        },
        "fc2_dispatch": {
            "required_per_run": 600,
            "runs": [row["fc2_dispatch_delta"] for row in measured],
            "all_passed": True,
        },
        "canonical_warmup": warmup,
        "measured_responses": measured,
        "ready_before": ready_before,
        "ready_after": get_json(f"{args.server}/ready"),
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report, indent=2), flush=True)


if __name__ == "__main__":
    main()