#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path
from typing import Any


SMOKE_WIDTH = 608
SMOKE_HEIGHT = 480
SMOKE_NUM_FRAMES = 124
SMOKE_SIGMA_POINTS = 50
MIN_REVIEWED_FRAMES = 4
SOURCE_MODEL_ID = "MiniMaxAI/MiniMax-H3"
SOURCE_REVISION = "73372e6cf53e414edd3ab03e357717fb0602e758"
LEARNED_COMPONENTS = frozenset({"transformer", "text_encoder", "vae", "audio_vae"})
BF16_ABLATION_COMPONENTS = frozenset({"transformer", "text_encoder"})


def timeline_frame_indices(frame_count: int, *, sample_count: int = 5) -> list[int]:
    if sample_count < 2:
        raise ValueError("timeline needs at least two samples")
    if frame_count < sample_count:
        raise ValueError("timeline needs at least as many frames as samples")
    last_index = frame_count - 1
    return [round(index * last_index / (sample_count - 1)) for index in range(sample_count)]


def build_component_plan(
    release: Path,
    *,
    task: str,
    bf16_components: set[str],
    component_paths: dict[str, Path] | None = None,
) -> dict[str, dict[str, str]]:
    component_paths = {} if component_paths is None else dict(component_paths)
    forbidden_vaes = set(bf16_components) & {"vae", "audio_vae"}
    if forbidden_vaes:
        raise ValueError(
            "generation runner always uses the source-precision VAE copies: "
            + ", ".join(sorted(forbidden_vaes))
        )
    unknown = set(bf16_components) - BF16_ABLATION_COMPONENTS
    if unknown:
        raise ValueError(f"unknown BF16 component: {', '.join(sorted(unknown))}")
    if task not in {"t2va", "ref2va"}:
        raise ValueError(f"unknown task: {task}")
    unknown_paths = set(component_paths) - LEARNED_COMPONENTS
    if unknown_paths:
        raise ValueError(f"unknown component path: {', '.join(sorted(unknown_paths))}")
    conflicts = set(bf16_components) & set(component_paths)
    if conflicts:
        raise ValueError(
            f"component cannot use both BF16 and a candidate path: {', '.join(sorted(conflicts))}"
        )

    transformer_subfolder = "transformer_ref" if task == "ref2va" else "transformer"
    subfolders = {
        "transformer": transformer_subfolder,
        "text_encoder": "text_encoder",
        "vae": "vae",
        "audio_vae": "audio_vae",
    }
    plan = {}
    for component, subfolder in subfolders.items():
        if component in bf16_components:
            plan[component] = {
                "source": "bf16",
                "model_id": SOURCE_MODEL_ID,
                "revision": SOURCE_REVISION,
                "subfolder": subfolder,
            }
        else:
            plan[component] = {
                "source": "release",
                "path": str(component_paths.get(component, release / subfolder)),
            }
    return plan


def evaluate_quality_smoke(metrics: dict[str, Any], review: dict[str, Any]) -> dict[str, Any]:
    width = int(metrics.get("width", 0))
    height = int(metrics.get("height", 0))
    if (width, height) != (SMOKE_WIDTH, SMOKE_HEIGHT):
        raise ValueError(f"quality smoke must be 608x480, got {width}x{height}")

    sigma_points = int(metrics.get("num_inference_steps", 0))
    if sigma_points != SMOKE_SIGMA_POINTS:
        raise ValueError(
            f"quality smoke must use {SMOKE_SIGMA_POINTS} sigma grid points, got {sigma_points}"
        )

    requested_frames = int(metrics.get("num_frames", 0))
    output_frames = int(metrics.get("video_frames", 0))
    if requested_frames != SMOKE_NUM_FRAMES or output_frames != SMOKE_NUM_FRAMES:
        raise ValueError(
            f"quality smoke must request and produce {SMOKE_NUM_FRAMES} frames, "
            f"got requested={requested_frames}, produced={output_frames}"
        )

    reviewed_frames = _validate_visual_review(review)

    return {
        "status": "pass",
        "model_evaluations": sigma_points - 1,
        "reviewed_frame_count": len(set(reviewed_frames)),
    }


def _validate_visual_review(review: dict[str, Any]) -> list[int]:
    if review.get("status") != "pass":
        raise ValueError("manual review must explicitly pass")
    reviewed_frames = list(review.get("reviewed_frame_indices", ()))
    if len(set(reviewed_frames)) < MIN_REVIEWED_FRAMES:
        raise ValueError(f"manual review must inspect at least {MIN_REVIEWED_FRAMES} distinct frames")
    if not review.get("prompt_subject_recognizable"):
        raise ValueError("manual review did not recognize the prompt subject")
    if not review.get("coherent_motion"):
        raise ValueError("manual review did not observe coherent motion")
    if review.get("repeating_tile_artifacts"):
        raise ValueError("manual review found repeating tile artifacts")
    if review.get("motion_ghosting") is not False:
        raise ValueError("manual review found motion ghosting or did not explicitly reject it")
    if review.get("texture_breakup") is not False:
        raise ValueError("manual review found texture breakup or did not explicitly reject it")
    if review.get("face_integrity") is not True:
        raise ValueError("manual review did not explicitly confirm face integrity")
    if review.get("full_resolution_frame_reviewed") is not True:
        raise ValueError("manual review did not inspect a full-resolution frame")
    return reviewed_frames