#!/usr/bin/env python3

import argparse
from pathlib import Path

import imageio.v2 as imageio


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Validate an indexed Self-Forcing video evaluation set."
    )
    parser.add_argument("video_dir", type=Path)
    parser.add_argument("--expected-videos", type=int, default=100)
    parser.add_argument("--expected-frames", type=int, default=81)
    parser.add_argument("--expected-height", type=int, default=480)
    parser.add_argument("--expected-width", type=int, default=832)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.video_dir.is_dir():
        raise FileNotFoundError(f"Video directory does not exist: {args.video_dir}")

    expected_names = {
        f"{index}-0_ema.mp4" for index in range(args.expected_videos)
    }
    video_paths = sorted(args.video_dir.glob("*.mp4"))
    actual_names = {path.name for path in video_paths}
    if actual_names != expected_names:
        missing = sorted(expected_names - actual_names)
        extra = sorted(actual_names - expected_names)
        raise ValueError(
            f"Video set mismatch: count={len(video_paths)}, "
            f"missing={missing[:10]}, extra={extra[:10]}"
        )

    for path in video_paths:
        reader = imageio.get_reader(path, format="ffmpeg")
        try:
            frame_count = reader.count_frames()
            if frame_count != args.expected_frames:
                raise ValueError(
                    f"{path.name}: expected {args.expected_frames} frames, "
                    f"found {frame_count}"
                )
            first_frame = reader.get_data(0)
            last_frame = reader.get_data(args.expected_frames - 1)
        finally:
            reader.close()
        expected_shape = (args.expected_height, args.expected_width, 3)
        if first_frame.shape != expected_shape or last_frame.shape != expected_shape:
            raise ValueError(
                f"{path.name}: expected frame shape {expected_shape}, "
                f"found first={first_frame.shape}, last={last_frame.shape}"
            )

    print(
        f"valid video set: path={args.video_dir.resolve()} "
        f"videos={len(video_paths)} frames_per_video={args.expected_frames} "
        f"resolution={args.expected_width}x{args.expected_height}"
    )


if __name__ == "__main__":
    main()