#!/usr/bin/env python3
"""Prepare EgoSpeak's Ego4D split as exact 2 FPS audio-video clips.

The 5 Hz target arrays are the timing authority.  For a target with T rows,
the prepared clip has duration T / 5 seconds and exactly ceil(2T / 5) video
frames.  Short source audio is padded with silence and all audio is trimmed to
the target duration.

Processing is intentionally sequential.  A single libx264 process uses all
available CPU cores on the intended extraction machine, so launching several
encoders normally creates CPU and disk contention instead of improving total
throughput.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np


TARGET_FPS = 5
VIDEO_FPS = 2
EXPECTED_TRAIN_CLIPS = 346
EXPECTED_TEST_CLIPS = 87


@dataclass(frozen=True)
class ClipSpec:
    split: str
    uid: str
    steps: int
    duration: float
    video_frames: int
    source: Path
    target: Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Prepare the 433 clips in EgoSpeak's EGO4DSHUFFLE split as "
            "strict fixed-grid 2 FPS MP4 files."
        )
    )
    parser.add_argument(
        "--data-root",
        type=Path,
        default=Path("/data/wangx/Datasets/ego4d"),
        help="Ego4D dataset root (default: %(default)s)",
    )
    parser.add_argument(
        "--split-json",
        type=Path,
        help="EgoSpeak LSTR data_info.json; defaults to DATA_ROOT/egospeak_data_info.json",
    )
    parser.add_argument(
        "--video-root",
        type=Path,
        help="Raw clip directory; defaults to DATA_ROOT/v1/clips",
    )
    parser.add_argument(
        "--target-root",
        type=Path,
        help="Target directory; defaults to DATA_ROOT/ego4d_target_frame",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        help="Prepared video directory; defaults to DATA_ROOT/omni_prepared_2fps",
    )
    parser.add_argument(
        "--manifest",
        type=Path,
        help="Output TSV; defaults to DATA_ROOT/ego4d_egospeak.tsv",
    )
    parser.add_argument(
        "--split",
        choices=("all", "train", "test"),
        default="all",
        help="Prepare all, train, or test clips (default: %(default)s)",
    )
    parser.add_argument(
        "--uid",
        action="append",
        help="Prepare only this split UID; repeat to select several clips",
    )
    parser.add_argument(
        "--limit",
        type=int,
        help="Prepare only the first N selected clips, for testing",
    )
    parser.add_argument(
        "--crf",
        type=int,
        default=18,
        help="libx264 constant-rate-factor value (default: %(default)s)",
    )
    parser.add_argument(
        "--preset",
        default="medium",
        help="libx264 preset (default: %(default)s)",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Rebuild outputs even when their frame counts are already correct",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validate inputs and print work without running ffmpeg",
    )
    args = parser.parse_args()

    if args.limit is not None and args.limit <= 0:
        parser.error("--limit must be positive")
    if not 0 <= args.crf <= 51:
        parser.error("--crf must be between 0 and 51")
    return args


def resolve_paths(args: argparse.Namespace) -> None:
    args.data_root = args.data_root.expanduser().resolve()
    args.split_json = (
        args.split_json or args.data_root / "egospeak_data_info.json"
    ).expanduser().resolve()
    args.video_root = (
        args.video_root or args.data_root / "v1" / "clips"
    ).expanduser().resolve()
    args.target_root = (
        args.target_root or args.data_root / "ego4d_target_frame"
    ).expanduser().resolve()
    args.output_root = (
        args.output_root or args.data_root / "omni_prepared_2fps"
    ).expanduser().resolve()
    args.manifest = (
        args.manifest or args.data_root / "ego4d_egospeak.tsv"
    ).expanduser().resolve()


def require_file(path: Path, description: str) -> None:
    if not path.is_file():
        raise FileNotFoundError(f"Missing {description}: {path}")


def load_specs(args: argparse.Namespace) -> list[ClipSpec]:
    require_file(args.split_json, "EgoSpeak split JSON")
    if not args.video_root.is_dir():
        raise FileNotFoundError(f"Missing raw-video directory: {args.video_root}")
    if not args.target_root.is_dir():
        raise FileNotFoundError(f"Missing target directory: {args.target_root}")

    with args.split_json.open() as file:
        payload = json.load(file)
    try:
        split_info = payload["EGO4DSHUFFLE"]
        train_uids = split_info["train_session_set"]
        test_uids = split_info["test_session_set"]
    except KeyError as error:
        raise RuntimeError(
            f"{args.split_json} is not EgoSpeak's LSTR data_info.json: "
            f"missing {error}"
        ) from error

    if len(train_uids) != EXPECTED_TRAIN_CLIPS:
        raise RuntimeError(
            f"Expected {EXPECTED_TRAIN_CLIPS} EGO4DSHUFFLE train clips, "
            f"found {len(train_uids)}"
        )
    if len(test_uids) != EXPECTED_TEST_CLIPS:
        raise RuntimeError(
            f"Expected {EXPECTED_TEST_CLIPS} EGO4DSHUFFLE test clips, "
            f"found {len(test_uids)}"
        )

    specs: list[ClipSpec] = []
    seen: set[str] = set()
    for split, uids in (("train", train_uids), ("test", test_uids)):
        for uid in uids:
            if uid in seen:
                raise RuntimeError(f"Duplicate UID in EGO4DSHUFFLE: {uid}")
            seen.add(uid)

            source = args.video_root / f"{uid}.mp4"
            target_path = args.target_root / f"{uid}.npy"
            require_file(source, "source video")
            require_file(target_path, "target array")

            target = np.load(target_path, mmap_mode="r", allow_pickle=False)
            if target.ndim != 2 or target.shape[1] != 3:
                raise RuntimeError(
                    f"Expected target shape [T, 3], found {target.shape}: "
                    f"{target_path}"
                )

            steps = int(target.shape[0])
            duration = steps / TARGET_FPS
            video_frames = math.ceil(steps * VIDEO_FPS / TARGET_FPS)
            specs.append(
                ClipSpec(
                    split=split,
                    uid=uid,
                    steps=steps,
                    duration=duration,
                    video_frames=video_frames,
                    source=source,
                    target=target_path,
                )
            )
    return specs


def write_manifest(specs: list[ClipSpec], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".partial")
    with temporary.open("w", newline="") as file:
        writer = csv.writer(file, delimiter="\t")
        writer.writerow(
            ["split", "uid", "steps", "duration", "video_frames_2fps"]
        )
        for spec in specs:
            writer.writerow(
                [
                    spec.split,
                    spec.uid,
                    spec.steps,
                    f"{spec.duration:.6f}",
                    spec.video_frames,
                ]
            )
    temporary.replace(path)


def select_specs(args: argparse.Namespace, specs: list[ClipSpec]) -> list[ClipSpec]:
    selected = specs
    if args.split != "all":
        selected = [spec for spec in selected if spec.split == args.split]
    if args.uid:
        requested = set(args.uid)
        known = {spec.uid for spec in specs}
        missing = sorted(requested - known)
        if missing:
            raise RuntimeError(
                "Requested UID(s) are not in EGO4DSHUFFLE: " + ", ".join(missing)
            )
        selected = [spec for spec in selected if spec.uid in requested]
    if args.limit is not None:
        selected = selected[: args.limit]
    if not selected:
        raise RuntimeError("No clips selected")
    return selected


def count_video_frames(path: Path) -> int | None:
    if not path.is_file() or path.stat().st_size == 0:
        return None
    command = [
        "ffprobe",
        "-v",
        "error",
        "-select_streams",
        "v:0",
        "-count_frames",
        "-show_entries",
        "stream=nb_read_frames",
        "-of",
        "default=noprint_wrappers=1:nokey=1",
        str(path),
    ]
    completed = subprocess.run(command, text=True, capture_output=True)
    if completed.returncode != 0:
        return None
    try:
        return int(completed.stdout.strip())
    except ValueError:
        return None


def prepare_one(
    args: argparse.Namespace,
    spec: ClipSpec,
    index: int,
    total: int,
) -> str:
    output = args.output_root / f"{spec.uid}_2fps.mp4"
    actual = count_video_frames(output)
    prefix = f"[{index}/{total}] [{spec.split}] {spec.uid}"

    if not args.force and actual == spec.video_frames:
        print(f"{prefix}: reuse ({actual} frames)", flush=True)
        return "reused"

    if actual is None:
        reason = "missing"
    else:
        reason = f"actual={actual}, expected={spec.video_frames}"
    print(
        f"{prefix}: prepare ({reason}; {spec.steps} states, "
        f"{spec.duration:.1f}s, {spec.video_frames} frames)",
        flush=True,
    )
    if args.dry_run:
        return "dry-run"

    temporary = args.output_root / f".{spec.uid}.partial.mp4"
    duration = f"{spec.duration:.6f}"
    video_filter = (
        f"[0:v:0]trim=start=0:end={duration},setpts=PTS-STARTPTS,"
        f"fps=fps={VIDEO_FPS}:start_time=0:round=down,"
        "tpad=stop_mode=clone:stop=2,"
        f"trim=start_frame=0:end_frame={spec.video_frames},"
        f"setpts=N/({VIDEO_FPS}*TB),format=yuv420p[vout]"
    )
    audio_filter = (
        f"[0:a:0]apad,atrim=start=0:end={duration},"
        "asetpts=PTS-STARTPTS[aout]"
    )
    command = [
        "ffmpeg",
        "-hide_banner",
        "-loglevel",
        "warning",
        "-stats",
        "-y",
        "-i",
        str(spec.source),
        "-filter_complex",
        f"{video_filter};{audio_filter}",
        "-map",
        "[vout]",
        "-map",
        "[aout]",
        "-map_metadata",
        "-1",
        "-c:v",
        "libx264",
        "-preset",
        args.preset,
        "-crf",
        str(args.crf),
        "-c:a",
        "aac",
        "-ar",
        "16000",
        "-ac",
        "1",
        "-b:a",
        "64k",
        "-movflags",
        "+faststart",
        str(temporary),
    ]
    completed = subprocess.run(command)
    if completed.returncode != 0:
        raise RuntimeError(
            f"ffmpeg failed for {spec.uid} with exit code {completed.returncode}; "
            f"partial output was left at {temporary}"
        )

    actual = count_video_frames(temporary)
    if actual != spec.video_frames:
        raise RuntimeError(
            f"Prepared frame count is {actual}, expected {spec.video_frames}: "
            f"{temporary}"
        )
    temporary.replace(output)
    print(f"{prefix}: OK ({actual} frames) -> {output}", flush=True)
    return "prepared"


def main() -> None:
    args = parse_args()
    resolve_paths(args)
    for executable in ("ffmpeg", "ffprobe"):
        if shutil.which(executable) is None:
            raise RuntimeError(f"Required executable is not on PATH: {executable}")

    specs = load_specs(args)
    write_manifest(specs, args.manifest)
    selected = select_specs(args, specs)
    args.output_root.mkdir(parents=True, exist_ok=True)

    total_steps = sum(spec.steps for spec in specs)
    print(
        f"Validated EGO4DSHUFFLE: {len(specs)} clips "
        f"({EXPECTED_TRAIN_CLIPS} train, {EXPECTED_TEST_CLIPS} test), "
        f"{total_steps} target states"
    )
    print(f"Manifest: {args.manifest}")
    print(f"Output:   {args.output_root}")
    print(f"Selected: {len(selected)} clips; sequential ffmpeg execution")

    counts = {"prepared": 0, "reused": 0, "dry-run": 0}
    for index, spec in enumerate(selected, start=1):
        result = prepare_one(args, spec, index, len(selected))
        counts[result] += 1

    print(
        "Complete: "
        + ", ".join(f"{key}={value}" for key, value in counts.items() if value)
    )


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nInterrupted; rerun the same command to resume.", file=sys.stderr)
        raise SystemExit(130)
    except Exception as error:
        print(f"ERROR: {error}", file=sys.stderr)
        raise SystemExit(1)