"""Delete extracted videos not referenced by sft_5k.parquet + eval_100.parquet.
Also delete raw archive files (zip/tar.gz) — annotations kept.

Idempotent. Prints what it would delete with --dry-run.
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

import pandas as pd

ROOT = Path("/mnt/local-fast/opd_zt")
DATA = ROOT / "data"
VIDEOS = DATA / "videos"
RAW = DATA / "raw"

VIDEO_EXTS = {".mp4", ".mkv", ".webm", ".avi", ".mov", ".m4v"}
ARCHIVE_EXTS = {".zip", ".tar", ".gz"}  # match .tar.gz via .gz too


def collect_keep_paths() -> set[Path]:
    """Read the 5K + eval parquets, return abs paths of every video referenced."""
    keep: set[Path] = set()
    for pq in [DATA / "sft_5k.parquet", DATA / "eval_100.parquet"]:
        if not pq.exists():
            print(f"WARN  missing parquet: {pq}", file=sys.stderr)
            continue
        df = pd.read_parquet(pq)
        for row in df["videos"]:
            for v in row:
                p = v["video"]
                if p.startswith("file://"):
                    p = p[len("file://"):]
                keep.add(Path(p).resolve())
    return keep


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--dry-run", action="store_true")
    p.add_argument("--keep-archives", action="store_true",
                   help="Do not delete .zip / .tar.gz under data/raw/")
    args = p.parse_args()

    keep = collect_keep_paths()
    print(f"[keep] {len(keep)} videos referenced by parquets")

    # --- 1. delete unused extracted videos ---
    n_kept = n_del = bytes_del = 0
    for root in [VIDEOS / "tb", VIDEOS / "lv178k"]:
        if not root.exists():
            continue
        for pp in root.rglob("*"):
            if pp.is_file() and pp.suffix.lower() in VIDEO_EXTS:
                if pp.resolve() in keep:
                    n_kept += 1
                else:
                    bytes_del += pp.stat().st_size
                    n_del += 1
                    if not args.dry_run:
                        pp.unlink()
    print(f"[videos] kept={n_kept}  deleted={n_del}  freed={bytes_del/1e9:.1f} GB")

    # --- 2. delete archives under data/raw/ ---
    if not args.keep_archives:
        n_arc = bytes_arc = 0
        for root in [RAW / "tb", RAW / "lv178k"]:
            if not root.exists():
                continue
            for pp in root.rglob("*"):
                if pp.is_file() and (
                    pp.suffix.lower() in ARCHIVE_EXTS
                    or pp.name.endswith(".tar.gz")
                    or pp.name.endswith(".tgz")
                ):
                    bytes_arc += pp.stat().st_size
                    n_arc += 1
                    if not args.dry_run:
                        pp.unlink()
        print(f"[archives] deleted={n_arc}  freed={bytes_arc/1e9:.1f} GB")

    # --- 3. prune empty directories ---
    if not args.dry_run:
        for root in [VIDEOS, RAW]:
            for pp in sorted(root.rglob("*"), key=lambda p: -len(p.parts)):
                if pp.is_dir():
                    try:
                        pp.rmdir()
                    except OSError:
                        pass

    print("DONE" + ("  [DRY-RUN]" if args.dry_run else ""))


if __name__ == "__main__":
    main()