#!/usr/bin/env python3
"""Build WorldMem train/test CSV index files from mp4/npz pairs.

Each row: video_path, action_path (paths relative to the project cwd when possible).

When --root is a symlink under the project (e.g. ./data), CSV rows use paths like
data/training/.... Avoid resolving symlinks so relative paths stay valid.

Example:
    python scripts/build_worldmem_csv.py
"""

from __future__ import annotations

import argparse
import csv
import random
from pathlib import Path


def _path_for_csv(path: Path) -> str:
    """Path for CSV rows, relative to cwd when possible (preserves symlinks)."""
    if not path.is_absolute():
        return path.as_posix()
    try:
        return path.relative_to(Path.cwd()).as_posix()
    except ValueError:
        return path.as_posix()


def collect_pairs(split_dir: Path, root: Path) -> list[tuple[str, str]]:
    """Collect (video_path, action_path) for every mp4 with a matching npz."""
    if not split_dir.is_dir():
        raise FileNotFoundError(f"Split directory not found: {split_dir}")

    pairs: list[tuple[str, str]] = []
    missing_npz: list[str] = []

    for mp4 in sorted(split_dir.rglob("*.mp4")):
        npz = mp4.with_suffix(".npz")
        if not npz.is_file():
            missing_npz.append(mp4.relative_to(root).as_posix())
            continue
        pairs.append((
            _path_for_csv(mp4),
            _path_for_csv(npz),
        ))

    if missing_npz:
        print(f"  warning: {len(missing_npz)} mp4 files without matching npz "
              f"(first: {missing_npz[0]})")

    return pairs


def write_csv(rows: list[tuple[str, str]], output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["video_path", "action_path"])
        writer.writerows(rows)
    print(f"  wrote {len(rows)} rows -> {output_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Build WorldMem train/test CSV metadata.")
    parser.add_argument(
        "--root",
        type=str,
        default='./data',
        help=f"Dataset root directory",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default='./data',
        help="Directory for train.csv / test.csv (default: same as --root)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for shuffling rows",
    )
    args = parser.parse_args()

    root = Path(args.root)
    output_dir = Path(args.output_dir) if args.output_dir else root

    splits = {
        "train.csv": root / "training",
        "test.csv": root / "test",
    }

    print(f"Dataset root : {root} (-> {root.resolve()})")
    print(f"Output dir   : {output_dir}")
    print(f"Shuffle seed : {args.seed}")

    rng = random.Random(args.seed)

    for csv_name, split_dir in splits.items():
        print(f"\n[{csv_name}] scanning {split_dir.relative_to(root)}/")
        pairs = collect_pairs(split_dir, root)
        rng.shuffle(pairs)
        write_csv(pairs, output_dir / csv_name)

    print("\nDone.")


if __name__ == "__main__":
    main()