import argparse
import json
import os
import sys


def ensure_dummy_dataset(base_dir: str, num_entries: int = 20) -> str:
    try:
        import numpy as np
        import imageio.v2 as imageio
    except Exception as exc:
        raise RuntimeError("imageio and numpy are required to build the dummy dataset.") from exc

    os.makedirs(base_dir, exist_ok=True)
    video_path = os.path.join(base_dir, "dummy.mp4")
    json_path = os.path.join(base_dir, "dummy_dataset.json")

    if not os.path.exists(video_path):
        writer = imageio.get_writer(video_path, fps=16)
        for i in range(90):
            frame = np.zeros((128, 128, 3), dtype=np.uint8)
            frame[:, :, 0] = (i * 3) % 255
            frame[:, :, 1] = (i * 7) % 255
            frame[:, :, 2] = (i * 13) % 255
            writer.append_data(frame)
        writer.close()

    facedetect_v1 = [
        [{"angle": {"yaw": 0, "pitch": 0, "roll": 0}, "detect": {"top": 10, "height": 100, "width": 100, "left": 10, "prob": 1.0}}],
        [{"angle": {"yaw": 60, "pitch": 0, "roll": 0}, "detect": {"top": 10, "height": 100, "width": 100, "left": 10, "prob": 1.0}}],
        [{"angle": {"yaw": -60, "pitch": 0, "roll": 0}, "detect": {"top": 10, "height": 100, "width": 100, "left": 10, "prob": 1.0}}],
    ]
    facedetect_v1_frame_index = [0, 20, 40]

    context = {
        "disk_path": video_path,
        "text": "dummy caption",
        "facedetect_v1": facedetect_v1,
        "facedetect_v1_frame_index": facedetect_v1_frame_index,
    }
    meta = {f"video_{i}": context for i in range(num_entries)}

    with open(json_path, "w") as f:
        json.dump(meta, f)

    return json_path


def main() -> int:
    parser = argparse.ArgumentParser(description="Multi-shot dataset dry-run.")
    parser.add_argument("--json_path", type=str, default="", help="Existing dataset JSON path.")
    parser.add_argument("--resolution", type=int, nargs=2, default=[64, 64], help="Resolution as H W.")
    parser.add_argument("--ref_num", type=int, default=3, help="Reference image count.")
    parser.add_argument("--base_dir", type=str, default="/data/rczhang/PencilFolder/multi-shot/tmp_dryrun")
    args = parser.parse_args()

    if args.json_path:
        json_path = args.json_path
    else:
        json_path = ensure_dummy_dataset(args.base_dir)

    sys.path.insert(0, "/data/rczhang/PencilFolder/multi-shot")
    from multi_view.datasets.videodataset import MulltiShot_MultiView_Dataset

    dataset = MulltiShot_MultiView_Dataset(
        dataset_base_path=json_path,
        resolution=tuple(args.resolution),
        ref_num=args.ref_num,
        training=True,
    )

    print("Dataset len:", len(dataset))
    item = dataset[0]
    print("Keys:", list(item.keys()))
    print("Video frames:", len(item["video"]))
    print("Ref images group:", len(item["ref_images"]), "x", len(item["ref_images"][0]))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())