# -*- coding: utf-8 -*-

import json
import os
import argparse

def count_frames(frames_dir: str, video_id: str) -> int:
    vid_dir = os.path.join(frames_dir, video_id)
    if not os.path.isdir(vid_dir):
        return 0
    return len([f for f in os.listdir(vid_dir) if f.endswith(".jpg")])

def main():
    p = argparse.ArgumentParser(
        description="Split test_videoids.json into short/long by frame count."
    )
    p.add_argument("--test_json",   required=True,
                   help="Path to test_videoids.json")
    p.add_argument("--frames_dir",  required=True,
                   help="Path to Video_Frames/ folder")
    p.add_argument("--threshold",   type=int, default=240,
                   help="Frame count threshold (default: 240)")
    p.add_argument("--out_dir",     required=True,
                   help="Output directory for split JSON files")
    args = p.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    with open(args.test_json, "r") as f:
        test_ids = json.load(f)

    short_ids, long_ids, missing_ids = [], [], []

    for vid in test_ids:
        n = count_frames(args.frames_dir, vid)
        if n == 0:
            missing_ids.append(vid)
            print(f"  [WARN] {vid}: no frames found — skipped")
        elif n < args.threshold:
            short_ids.append(vid)
            print(f"  [SHORT] {vid}: {n} frames")
        else:
            long_ids.append(vid)
            print(f"  [LONG]  {vid}: {n} frames")

    short_path = os.path.join(args.out_dir, "test_videoids_short.json")
    long_path  = os.path.join(args.out_dir, "test_videoids_long.json")

    with open(short_path, "w") as f:
        json.dump(short_ids, f, indent=2)
    with open(long_path, "w") as f:
        json.dump(long_ids, f, indent=2)

    print(f"\n  Total test videos : {len(test_ids)}")
    print(f"  Short (<{args.threshold} frames) : {len(short_ids)} -> {short_path}")
    print(f"  Long  (>={args.threshold} frames): {len(long_ids)} -> {long_path}")
    if missing_ids:
        print(f"  Missing (no frames): {len(missing_ids)} -> {missing_ids}")

if __name__ == "__main__":
    main()