import json
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import decord

CACHE = Path("/home/s56fasla/MOTIONX/cache")
WORKERS = 16


def probe(item):
    try:
        vr = decord.VideoReader(item["video"], num_threads=1)
        n = len(vr)
        if n < 2:
            return item, "too_short"
        vr[0]
        return item, None
    except Exception as e:
        return item, type(e).__name__


for split in ["train", "val"]:
    items = json.load(open(CACHE / f"manifest_{split}.json"))
    bad = []
    good = []
    with ThreadPoolExecutor(WORKERS) as ex:
        for item, err in ex.map(probe, items):
            if err:
                bad.append({"video": item["video"], "err": err})
            else:
                good.append(item)
    json.dump(good, open(CACHE / f"manifest_{split}_ok.json", "w"))
    json.dump(bad, open(CACHE / f"bad_videos_{split}.json", "w"))
    print(split, "ok", len(good), "bad", len(bad), flush=True)
    for b in bad[:5]:
        print("   ", b["err"], Path(b["video"]).name, flush=True)