#!/usr/bin/env python3
"""Select representative files and dependency-forcing mid-GOP edit ranges."""

from __future__ import annotations

import argparse
from dataclasses import dataclass
import json
from pathlib import Path
import subprocess
import sys
from typing import Any


@dataclass(frozen=True)
class Target:
    name: str
    extension: str
    video: str
    depth: int
    audio: str


TARGETS = [
    Target("mpeg1-mpg-mp2", ".mpg", "mpeg1video", 8, "mp2"),
    Target("mpeg2-mpg-mp2", ".mpg", "mpeg2video", 8, "mp2"),
    Target("mpeg2-mpg-ac3", ".mpg", "mpeg2video", 8, "ac3"),
    Target("mpeg2-ts-ac3", ".ts", "mpeg2video", 8, "ac3"),
    Target("mpeg2-m2ts-mp2", ".m2ts", "mpeg2video", 8, "mp2"),
    Target("mpeg2-mkv-ac3", ".mkv", "mpeg2video", 8, "ac3"),
    Target("h264-mp4-aac", ".mp4", "h264", 8, "aac"),
    Target("h264-mov-aac", ".mov", "h264", 8, "aac"),
    Target("h264-m2ts-ac3", ".m2ts", "h264", 8, "ac3"),
    Target("h264-ts-ac3", ".ts", "h264", 8, "ac3"),
    Target("h264-mkv-aac", ".mkv", "h264", 8, "aac"),
    Target("h264-mkv-ac3", ".mkv", "h264", 8, "ac3"),
    Target("h264-mkv-eac3", ".mkv", "h264", 8, "eac3"),
    Target("h264-mkv-dts", ".mkv", "h264", 8, "dts"),
    Target("h264-mkv-flac", ".mkv", "h264", 8, "flac"),
    Target("h264-mkv-mp3", ".mkv", "h264", 8, "mp3"),
    Target("h264-mov-pcm", ".mov", "h264", 8, "pcm_s16le"),
    Target("h264-main10-mkv-aac", ".mkv", "h264", 10, "aac"),
    Target("hevc-main-mp4-aac", ".mp4", "hevc", 8, "aac"),
    Target("hevc-main-mkv-aac", ".mkv", "hevc", 8, "aac"),
    Target("hevc-main-mkv-ac3", ".mkv", "hevc", 8, "ac3"),
    Target("hevc-main-mkv-eac3", ".mkv", "hevc", 8, "eac3"),
    Target("hevc-main10-mp4-aac", ".mp4", "hevc", 10, "aac"),
    Target("hevc-main10-mkv-aac", ".mkv", "hevc", 10, "aac"),
    Target("hevc-main10-mkv-ac3", ".mkv", "hevc", 10, "ac3"),
    Target("hevc-main10-mkv-eac3", ".mkv", "hevc", 10, "eac3"),
    Target("hevc-main10-mkv-dts", ".mkv", "hevc", 10, "dts"),
    Target("hevc-main10-mkv-truehd", ".mkv", "hevc", 10, "truehd"),
    Target("vp9-mkv-opus", ".mkv", "vp9", 8, "opus"),
    Target("vp9-mp4-aac", ".mp4", "vp9", 8, "aac"),
    Target("av1-main10-mkv-opus", ".mkv", "av1", 10, "opus"),
    Target("av1-main10-mkv-eac3", ".mkv", "av1", 10, "eac3"),
    Target("mpeg4-avi-mp3", ".avi", "mpeg4", 8, "mp3"),
    Target("vc1-m2ts-ac3", ".m2ts", "vc1", 8, "ac3"),
    Target("dv-avi-pcm", ".avi", "dvvideo", 8, "pcm_s16le"),
    Target("wmv3-wmv-wmav2", ".wmv", "wmv3", 8, "wmav2"),
]


def output_profile(extension: str) -> tuple[str, str]:
    if extension in {".mp4", ".m4v", ".mov"}:
        return "MP4", ".mp4"
    if extension == ".ts":
        return "Transport Stream", ".ts"
    if extension in {".m2ts", ".mts"}:
        return "M2TS", ".m2ts"
    if extension in {".mpg", ".mpeg", ".vob"}:
        return "Program Stream", ".mpg"
    return "Matroska MKV", ".mkv"


def load_items(paths: list[Path]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for priority, path in enumerate(paths):
        document = json.loads(path.read_text(encoding="utf-8"))
        for row in document["items"]:
            row = dict(row)
            row["inventory_priority"] = priority
            rows.append(row)
    return rows


def candidate_rows(rows: list[dict[str, Any]], target: Target) -> list[dict[str, Any]]:
    matches = []
    for row in rows:
        if row.get("status") != "ok" or row["extension"] != target.extension:
            continue
        if row["video"]["codec"] != target.video or row["video"]["bit_depth"] != target.depth:
            continue
        if target.audio not in {audio["codec"] for audio in row["audio"]}:
            continue
        if row["duration"] < 8:
            continue
        matches.append(row)
    return sorted(
        matches,
        key=lambda row: (
            row["inventory_priority"],
            row["video"]["width"] > 1920 or row["video"]["height"] > 1080,
            row["size"],
        ),
    )


def indexed_pictures(ffprobe: str, path: str, video_codec: str) -> list[dict[str, Any]]:
    decode_picture_types = video_codec in {"mpeg1video", "mpeg2video"}
    interval = "%+45" if decode_picture_types else "%+180"
    command = [ffprobe, "-v", "error", "-select_streams", "v:0", "-read_intervals", interval]
    if decode_picture_types:
        command += ["-show_frames", "-show_entries", "frame=key_frame,pict_type,best_effort_timestamp_time"]
    else:
        command += ["-show_packets", "-show_entries", "packet=pts_time,flags"]
    command += ["-of", "json", path]
    try:
        completed = subprocess.run(command, capture_output=True, text=True, timeout=60, encoding="utf-8", errors="replace")
    except subprocess.TimeoutExpired:
        return []
    if completed.returncode:
        return []
    document = json.loads(completed.stdout)
    if decode_picture_types:
        return document.get("frames", [])
    packets = [packet for packet in document.get("packets", []) if "pts_time" in packet]
    packets.sort(key=lambda packet: float(packet["pts_time"]))
    return [
        {
            "best_effort_timestamp_time": packet["pts_time"],
            "key_frame": int("K" in packet.get("flags", "")),
            "pict_type": "",
        }
        for packet in packets
    ]


def choose_boundaries(frames: list[dict[str, Any]]) -> dict[str, Any] | None:
    usable = [frame for frame in frames if "best_effort_timestamp_time" in frame]
    keys = [index for index, frame in enumerate(usable) if int(frame.get("key_frame", 0))]
    if len(keys) < 4:
        return None
    for first_key in range(len(keys) - 3):
        start_gop_begin, start_gop_end = keys[first_key], keys[first_key + 1]
        tail_gop_begin, tail_gop_end = keys[first_key + 2], keys[first_key + 3]
        starts = [
            index for index in range(start_gop_begin + 2, start_gop_end - 1)
            if not int(usable[index].get("key_frame", 0))
        ]
        ends = [
            index for index in range(tail_gop_begin + 2, tail_gop_end - 1)
            if not int(usable[index].get("key_frame", 0))
            and usable[index].get("pict_type") == "B"
        ]
        if not ends:
            ends = [
                index for index in range(tail_gop_begin + 2, tail_gop_end - 1)
                if not int(usable[index].get("key_frame", 0))
            ]
        if not starts or not ends:
            continue
        start = starts[min(2, len(starts) - 1)]
        end = ends[min(2, len(ends) - 1)]
        if end <= start + 2:
            continue
        start_ms = round(float(usable[start]["best_effort_timestamp_time"]) * 1000)
        end_ms = round(float(usable[end]["best_effort_timestamp_time"]) * 1000)
        end_exclusive_ms = round(float(usable[end + 1]["best_effort_timestamp_time"]) * 1000)
        return {
            "start_frame": start,
            "end_frame_inclusive": end,
            "end_frame_exclusive": end + 1,
            "start_ms": start_ms,
            "end_ms": end_ms,
            "end_exclusive_ms": end_exclusive_ms,
            "start_picture": usable[start].get("pict_type", ""),
            "end_picture": usable[end].get("pict_type", ""),
            "keyframes": keys[first_key:first_key + 4],
            "expected_frames": end - start + 1,
        }
    return None


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--inventory", action="append", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--ffprobe", default=r"C:\ffmpeg_latest\ffprobe.exe")
    arguments = parser.parse_args()
    items = load_items(arguments.inventory)
    cases: list[dict[str, Any]] = []
    missing: list[str] = []
    for target in TARGETS:
        selected = None
        for candidate in candidate_rows(items, target)[:8]:
            boundaries = choose_boundaries(indexed_pictures(arguments.ffprobe, candidate["path"], target.video))
            if boundaries:
                selected = (candidate, boundaries)
                break
        if not selected:
            missing.append(target.name)
            print(f"missing {target.name}", flush=True)
            continue
        candidate, boundaries = selected
        profile, output_extension = output_profile(target.extension)
        builds = ["original", "updated"] if target.depth > 8 else ["original"]
        case = {
            "name": target.name,
            "source": candidate["path"],
            "source_bytes": candidate["size"],
            "source_extension": target.extension,
            "source_format": candidate["format"],
            "video": candidate["video"],
            "audio": candidate["audio"],
            "stream_count": candidate["stream_count"],
            "output_profile": profile,
            "output_extension": output_extension,
            "videoredo_builds": builds,
            **boundaries,
        }
        cases.append(case)
        print(f"selected {target.name}: {candidate['path']} [{boundaries['start_frame']}..{boundaries['end_frame_inclusive']}]", flush=True)
    document = {
        "schema": "smartcut-videoredo-comparison-manifest-v1",
        "inventories": [str(path) for path in arguments.inventory],
        "cases": cases,
        "missing": missing,
    }
    arguments.output.parent.mkdir(parents=True, exist_ok=True)
    temporary = arguments.output.with_suffix(arguments.output.suffix + ".tmp")
    temporary.write_text(json.dumps(document, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    temporary.replace(arguments.output)
    print(f"wrote {arguments.output}: {len(cases)} cases, {len(missing)} missing")
    return 0


if __name__ == "__main__":
    sys.exit(main())
