#!/usr/bin/env python3
"""Compute paired, frame-aligned YUV420 SSIM for MiniMax-H3 runs.

The candidate and reference directories must each contain the throughput
client ``results.jsonl`` files. Rows are paired by ``request_id`` and checked
for matching prompt, seed, task, geometry, duration, and aspect ratio before
the decoded videos are compared with FFmpeg's native ``ssim`` filter.
"""

from __future__ import annotations

import argparse
import json
import math
import shutil
import statistics
import subprocess
import tempfile
from collections import defaultdict
from pathlib import Path
from typing import Any


PAIR_FIELDS = (
    "task",
    "short_edge",
    "prompt_index",
    "prompt",
    "seed",
    "duration_seconds",
    "aspect_ratio",
)


def percentile(values: list[float], q: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    pos = (len(ordered) - 1) * q
    lo, hi = math.floor(pos), math.ceil(pos)
    if lo == hi:
        return ordered[lo]
    return ordered[lo] * (hi - pos) + ordered[hi] * (pos - lo)


def resolve_executable(explicit: str | None, name: str) -> str:
    if explicit:
        path = Path(explicit)
        if path.is_file():
            return str(path)
        resolved = shutil.which(explicit)
        if resolved:
            return resolved
        raise SystemExit(f"{name} executable not found: {explicit}")

    resolved = shutil.which(name)
    if resolved:
        return resolved
    candidates = [
        Path("/root/.miniconda3/envs/deploy/bin") / name,
        Path("/root/.miniconda3/envs/vllm/bin") / name,
        Path("/root/.miniconda3/envs/bbj/bin") / name,
    ]
    for path in candidates:
        if path.is_file():
            return str(path)
    raise SystemExit(f"{name} is required; pass --{name} explicitly")


def run_checked(command: list[str]) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(command, text=True, capture_output=True, check=False)
    if result.returncode:
        rendered = " ".join(command)
        raise RuntimeError(
            f"command failed ({result.returncode}): {rendered}\n{result.stderr[-4000:]}"
        )
    return result


def resolve_video_path(root: Path, raw_path: str) -> Path:
    if not raw_path:
        raise ValueError("file_path is empty")
    path = Path(raw_path)
    if path.is_file():
        return path.resolve()
    matches = [candidate for candidate in root.rglob(path.name) if candidate.is_file()]
    if len(matches) == 1:
        return matches[0].resolve()
    if not matches:
        raise ValueError(f"video does not exist: {path}")
    raise ValueError(
        f"video path {path} is stale and filename is ambiguous under {root}: "
        + ", ".join(str(match) for match in matches[:10])
    )


def load_rows(root: Path) -> dict[str, dict[str, Any]]:
    rows: dict[str, dict[str, Any]] = {}
    files = sorted(root.glob("client_*/results.jsonl"))
    if not files:
        files = sorted(root.rglob("client_*/results.jsonl"))
    if not files:
        raise ValueError(f"no client_*/results.jsonl found under {root}")

    for path in files:
        for line_number, line in enumerate(
            path.read_text(encoding="utf-8").splitlines(), start=1
        ):
            if not line.strip():
                continue
            row = json.loads(line)
            if not row.get("success"):
                continue
            request_id = str(row.get("request_id") or "")
            if not request_id:
                raise ValueError(f"missing request_id: {path}:{line_number}")
            if request_id in rows:
                raise ValueError(
                    f"duplicate request_id {request_id!r} under {root}; "
                    "pass one task/topology phase rather than a whole matrix"
                )
            try:
                video_path = resolve_video_path(root, str(row.get("file_path") or ""))
            except ValueError as error:
                raise ValueError(f"video for {request_id!r}: {error}") from error
            row["_resolved_file_path"] = str(video_path)
            rows[request_id] = row
    return rows


def check_pair(candidate: dict[str, Any], reference: dict[str, Any]) -> None:
    mismatches = []
    for field in PAIR_FIELDS:
        if candidate.get(field) != reference.get(field):
            mismatches.append(
                f"{field}: candidate={candidate.get(field)!r} "
                f"reference={reference.get(field)!r}"
            )
    if mismatches:
        raise ValueError("pair metadata mismatch: " + "; ".join(mismatches))


def probe_video(ffprobe: str, path: Path) -> dict[str, Any]:
    result = run_checked(
        [
            ffprobe,
            "-v",
            "error",
            "-select_streams",
            "v:0",
            "-count_frames",
            "-show_entries",
            "stream=width,height,pix_fmt,r_frame_rate,avg_frame_rate,nb_frames,nb_read_frames",
            "-of",
            "json",
            str(path),
        ]
    )
    payload = json.loads(result.stdout)
    streams = payload.get("streams") or []
    if len(streams) != 1:
        raise ValueError(f"expected one video stream in {path}, got {len(streams)}")
    stream = streams[0]
    frame_text = stream.get("nb_read_frames") or stream.get("nb_frames")
    if frame_text in (None, "N/A"):
        raise ValueError(f"could not determine decoded frame count for {path}")
    return {
        "width": int(stream["width"]),
        "height": int(stream["height"]),
        "pix_fmt": stream.get("pix_fmt"),
        "r_frame_rate": stream.get("r_frame_rate"),
        "avg_frame_rate": stream.get("avg_frame_rate"),
        "frames": int(frame_text),
    }


def check_video_contract(candidate: dict[str, Any], reference: dict[str, Any]) -> None:
    fields = ("width", "height", "r_frame_rate", "frames")
    mismatches = [
        f"{field}: candidate={candidate[field]!r} reference={reference[field]!r}"
        for field in fields
        if candidate[field] != reference[field]
    ]
    if mismatches:
        raise ValueError("decoded video mismatch: " + "; ".join(mismatches))


def parse_ffmpeg_stats(path: Path) -> list[dict[str, float | int]]:
    frames: list[dict[str, float | int]] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        values: dict[str, str] = {}
        for token in line.split():
            if ":" in token:
                key, value = token.split(":", 1)
                values[key] = value
        if not {"n", "Y", "U", "V", "All"}.issubset(values):
            continue
        frames.append(
            {
                "frame": int(values["n"]),
                "y": float(values["Y"]),
                "u": float(values["U"]),
                "v": float(values["V"]),
                "all": float(values["All"]),
            }
        )
    if not frames:
        raise ValueError(f"FFmpeg emitted no per-frame SSIM metrics: {path}")
    return frames


def compare_video_pair(
    ffmpeg: str,
    ffprobe: str,
    candidate_path: Path,
    reference_path: Path,
    stats_path: Path,
) -> tuple[dict[str, Any], list[dict[str, float | int]]]:
    candidate_probe = probe_video(ffprobe, candidate_path)
    reference_probe = probe_video(ffprobe, reference_path)
    check_video_contract(candidate_probe, reference_probe)

    filter_graph = (
        "[0:v]setpts=PTS-STARTPTS,format=yuv420p[candidate];"
        "[1:v]setpts=PTS-STARTPTS,format=yuv420p[reference];"
        f"[candidate][reference]ssim=stats_file={stats_path}"
    )
    run_checked(
        [
            ffmpeg,
            "-hide_banner",
            "-nostdin",
            "-loglevel",
            "error",
            "-i",
            str(candidate_path),
            "-i",
            str(reference_path),
            "-filter_complex",
            filter_graph,
            "-an",
            "-f",
            "null",
            "-",
        ]
    )
    frames = parse_ffmpeg_stats(stats_path)
    if len(frames) != candidate_probe["frames"]:
        raise ValueError(
            f"SSIM frame count mismatch: metrics={len(frames)} "
            f"decoded={candidate_probe['frames']}"
        )
    return candidate_probe, frames


def metric_summary(values: list[float]) -> dict[str, float]:
    return {
        "mean": statistics.fmean(values),
        "p10": percentile(values, 0.10),
        "min": min(values),
        "max": max(values),
    }


def compare_command(args: argparse.Namespace) -> int:
    if not 0.0 <= args.threshold <= 1.0:
        raise SystemExit("--threshold must be between 0 and 1")
    ffmpeg = resolve_executable(args.ffmpeg, "ffmpeg")
    ffprobe = resolve_executable(args.ffprobe, "ffprobe")
    candidates = load_rows(args.candidate_dir)
    references = load_rows(args.reference_dir)

    candidate_ids = set(candidates)
    reference_ids = set(references)
    if candidate_ids != reference_ids:
        missing_reference = sorted(candidate_ids - reference_ids)
        missing_candidate = sorted(reference_ids - candidate_ids)
        raise SystemExit(
            "request sets do not match: "
            f"missing_reference={missing_reference[:20]} "
            f"missing_candidate={missing_candidate[:20]}"
        )

    selected_ids = sorted(candidate_ids)
    if args.limit is not None:
        if args.limit < 1:
            raise SystemExit("--limit must be at least 1")
        selected_ids = selected_ids[: args.limit]

    args.output_dir.mkdir(parents=True, exist_ok=True)
    pair_rows: list[dict[str, Any]] = []
    all_frame_rows: list[dict[str, Any]] = []
    with tempfile.TemporaryDirectory(prefix="h3-paired-ssim-") as temporary:
        temporary_root = Path(temporary)
        for index, request_id in enumerate(selected_ids, start=1):
            candidate = candidates[request_id]
            reference = references[request_id]
            check_pair(candidate, reference)
            probe, frames = compare_video_pair(
                ffmpeg,
                ffprobe,
                Path(candidate["_resolved_file_path"]),
                Path(reference["_resolved_file_path"]),
                temporary_root / f"{index:05d}.stats",
            )
            all_values = [float(frame["all"]) for frame in frames]
            y_values = [float(frame["y"]) for frame in frames]
            u_values = [float(frame["u"]) for frame in frames]
            v_values = [float(frame["v"]) for frame in frames]
            all_summary = metric_summary(all_values)
            pair = {
                "request_id": request_id,
                "task": candidate["task"],
                "short_edge": candidate["short_edge"],
                "prompt_index": candidate["prompt_index"],
                "prompt": candidate["prompt"],
                "seed": candidate["seed"],
                "candidate_file": candidate["_resolved_file_path"],
                "reference_file": reference["_resolved_file_path"],
                "candidate_num_inference_steps": candidate.get("num_inference_steps"),
                "reference_num_inference_steps": reference.get("num_inference_steps"),
                "width": probe["width"],
                "height": probe["height"],
                "fps": probe["r_frame_rate"],
                "frames": len(frames),
                "ssim_all_mean": all_summary["mean"],
                "ssim_all_p10": all_summary["p10"],
                "ssim_all_min": all_summary["min"],
                "ssim_y_mean": statistics.fmean(y_values),
                "ssim_u_mean": statistics.fmean(u_values),
                "ssim_v_mean": statistics.fmean(v_values),
                "threshold": args.threshold,
                "passed": all_summary["mean"] >= args.threshold,
            }
            pair_rows.append(pair)
            for frame in frames:
                all_frame_rows.append({"request_id": request_id, **frame})
            print(
                f"[{index}/{len(selected_ids)}] {request_id} "
                f"mean={all_summary['mean']:.6f} p10={all_summary['p10']:.6f} "
                f"min={all_summary['min']:.6f}",
                flush=True,
            )

    all_values = [float(row["all"]) for row in all_frame_rows]
    video_means = [float(row["ssim_all_mean"]) for row in pair_rows]
    by_short_edge: dict[str, dict[str, Any]] = {}
    grouped: dict[int, list[float]] = defaultdict(list)
    for row in pair_rows:
        grouped[int(row["short_edge"])].append(float(row["ssim_all_mean"]))
    for short_edge, values in sorted(grouped.items()):
        by_short_edge[str(short_edge)] = {
            "videos": len(values),
            "mean_video_ssim": statistics.fmean(values),
            "min_video_ssim": min(values),
        }

    frame_summary = metric_summary(all_values)
    summary = {
        "metric": "FFmpeg decoded YUV420 SSIM All",
        "aggregation": {
            "mean_video_ssim": statistics.fmean(video_means),
            "frame_weighted_mean_ssim": frame_summary["mean"],
            "frame_p10_ssim": frame_summary["p10"],
            "min_frame_ssim": frame_summary["min"],
        },
        "candidate_dir": str(args.candidate_dir.resolve()),
        "reference_dir": str(args.reference_dir.resolve()),
        "ffmpeg": ffmpeg,
        "ffprobe": ffprobe,
        "threshold": args.threshold,
        "overall_passed": statistics.fmean(video_means) >= args.threshold,
        "videos": len(pair_rows),
        "videos_passed": sum(bool(row["passed"]) for row in pair_rows),
        "frames": len(all_frame_rows),
        "by_short_edge": by_short_edge,
        "pairs": pair_rows,
    }
    (args.output_dir / "paired_ssim.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )

    pair_columns = [
        "request_id",
        "task",
        "short_edge",
        "prompt_index",
        "seed",
        "width",
        "height",
        "fps",
        "frames",
        "ssim_all_mean",
        "ssim_all_p10",
        "ssim_all_min",
        "ssim_y_mean",
        "ssim_u_mean",
        "ssim_v_mean",
        "threshold",
        "passed",
        "candidate_file",
        "reference_file",
    ]
    with (args.output_dir / "paired_ssim.tsv").open("w", encoding="utf-8") as handle:
        handle.write("\t".join(pair_columns) + "\n")
        for row in pair_rows:
            handle.write("\t".join(str(row[column]) for column in pair_columns) + "\n")

    frame_columns = ("request_id", "frame", "y", "u", "v", "all")
    with (args.output_dir / "paired_ssim_frames.tsv").open(
        "w", encoding="utf-8"
    ) as handle:
        handle.write("\t".join(frame_columns) + "\n")
        for row in all_frame_rows:
            handle.write("\t".join(str(row[column]) for column in frame_columns) + "\n")

    print(json.dumps(summary["aggregation"], ensure_ascii=False), flush=True)
    if args.fail_below_threshold and not summary["overall_passed"]:
        return 2
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    compare = subparsers.add_parser("compare")
    compare.add_argument("--candidate-dir", type=Path, required=True)
    compare.add_argument("--reference-dir", type=Path, required=True)
    compare.add_argument("--output-dir", type=Path, required=True)
    compare.add_argument("--threshold", type=float, default=0.90)
    compare.add_argument(
        "--limit",
        type=int,
        help="compare only the first N matched requests (intended for smoke tests)",
    )
    compare.add_argument("--ffmpeg")
    compare.add_argument("--ffprobe")
    compare.add_argument("--fail-below-threshold", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.command == "compare":
        return compare_command(args)
    raise AssertionError(args.command)


if __name__ == "__main__":
    raise SystemExit(main())