"""
Video / frame loading for experiments.
"""
from __future__ import annotations
from pathlib import Path
from PIL import Image
from util.paths import VIDEO_DIR, get_frames_dir

def resolve_video_path(video_file_id: str) -> Path:
    """Return absolute path to the source mp4. Raises FileNotFoundError if missing."""
    p = VIDEO_DIR / f'{video_file_id}.mp4'
    if not p.exists():
        raise FileNotFoundError(
            f'Source video not found: {p}\n'
            f'  Set CHNB_VIDEO_DIR to the directory containing {video_file_id}.mp4.'
        )
    return p


def load_uniform_frames(video_file_id: str, n: int, *, strict: bool = True) -> list[Image.Image]:
    """
    Load N pre-extracted uniform frames for a video.

    Expects JPEGs at `data/frames/{video_file_id}/uniform_{n}/*.jpg` (named
    0000.jpg .. 000{n-1}.jpg). Produced by `experiments/preprocess/run_server.py`.

    If `strict=True` (default), raises FileNotFoundError when the directory
    doesn't exist or has fewer than N JPEGs. Set strict=False to return whatever
    is present (may be an empty list).
    """
    d = get_frames_dir(video_file_id, 'uniform', n)
    if not d.exists():
        if strict:
            raise FileNotFoundError(
                f'Pre-extracted frames missing: {d}\n'
                f'  Run: python experiments/preprocess/run_server.py '
                f'--uniform-n {n}'
            )
        return []

    jpgs = sorted(d.glob('*.jpg'))
    if strict and len(jpgs) < n:
        raise FileNotFoundError(
            f'Expected {n} frames at {d} but found only {len(jpgs)}.\n'
            f'  Re-run preprocess to regenerate this budget.'
        )

    images: list[Image.Image] = []
    for f in jpgs:
        img = Image.open(f)
        img.load()   # force read so the file handle can close
        images.append(img)
    return images