#!/usr/bin/env python3
"""Build VBench custom-input directories from dynamic-gate video outputs."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path


DIMENSIONS = [
    "subject_consistency",
    "background_consistency",
    "motion_smoothness",
    "aesthetic_quality",
    "imaging_quality",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset_root", type=Path, required=True)
    parser.add_argument("--output_root", type=Path, required=True)
    parser.add_argument(
        "--condition",
        action="append",
        nargs=2,
        metavar=("NAME", "VIDEO_DIR"),
        required=True,
        help="Condition name and directory containing prompt_XXXX.mp4 files.",
    )
    parser.add_argument("--prompt_start", type=int, default=90)
    parser.add_argument("--prompt_end", type=int, default=100)
    parser.add_argument(
        "--combined_name",
        default=None,
        help="Also build one flattened input directory for a single shared VBench run.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    dataset_root = args.dataset_root.resolve()
    output_root = args.output_root.resolve()
    output_root.mkdir(parents=True, exist_ok=True)
    combined_dir = output_root / args.combined_name if args.combined_name else None
    combined_info = []
    if combined_dir is not None:
        combined_dir.mkdir(parents=True, exist_ok=True)

    for name, source_text in args.condition:
        source_dir = Path(source_text).resolve()
        destination_dir = output_root / name
        destination_dir.mkdir(parents=True, exist_ok=True)
        full_info = []
        for prompt_id in range(args.prompt_start, args.prompt_end):
            filename = f"prompt_{prompt_id:04d}.mp4"
            source = source_dir / filename
            if not source.is_file():
                raise FileNotFoundError(source)
            destination = destination_dir / filename
            if destination.is_symlink() or destination.exists():
                destination.unlink()
            destination.symlink_to(os.path.relpath(source, destination_dir))

            metadata_path = dataset_root / f"prompt_{prompt_id:04d}" / "metadata.json"
            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
            full_info.append(
                {
                    "video_list": [filename],
                    "prompt_en": metadata["prompt"],
                    "dimension": DIMENSIONS,
                }
            )
            if combined_dir is not None:
                combined_filename = f"{name}__{filename}"
                combined_destination = combined_dir / combined_filename
                if combined_destination.is_symlink() or combined_destination.exists():
                    combined_destination.unlink()
                combined_destination.symlink_to(
                    os.path.relpath(source, combined_dir)
                )
                combined_info.append(
                    {
                        "video_list": [combined_filename],
                        "prompt_en": metadata["prompt"],
                        "dimension": DIMENSIONS,
                    }
                )
        (destination_dir / "full_info.json").write_text(
            json.dumps(full_info, indent=2) + "\n", encoding="utf-8"
        )
        print(f"[prepared] {name}: {len(full_info)} videos -> {destination_dir}")

    if combined_dir is not None:
        (combined_dir / "full_info.json").write_text(
            json.dumps(combined_info, indent=2) + "\n", encoding="utf-8"
        )
        print(
            f"[prepared] combined: {len(combined_info)} videos -> {combined_dir}"
        )


if __name__ == "__main__":
    main()