"""
Shared utilities for eval_metrics: discover evals_ep0 outputs and load video frames.
"""
from __future__ import annotations

import os
from typing import List, Tuple

try:
    import cv2
    HAS_CV2 = True
except ImportError:
    HAS_CV2 = False

import numpy as np


def discover_evals_videos(evals_root: str, pattern: str = "*_gen_only.mp4") -> List[Tuple[str, str]]:
    """
    Discover all generated-only MP4s under evals_ep0 structure.
    Returns list of (relative_path, absolute_path) for each video.
    """
    out: List[Tuple[str, str]] = []
    evals_root = os.path.abspath(evals_root)
    for root, _dirs, files in os.walk(evals_root):
        for f in files:
            if f.endswith("_gen_only.mp4") or (pattern != "*_gen_only.mp4" and f.endswith(".mp4")):
                absp = os.path.join(root, f)
                rel = os.path.relpath(absp, evals_root)
                out.append((rel, absp))
    return sorted(out, key=lambda x: x[0])


def discover_loop_closure_videos(evals_root: str) -> List[Tuple[str, str]]:
    """Discover MP4s under 1_loop_4chunk and 3_multi_ctx_4chunk for loop closure (prefer gen_only)."""
    out: List[Tuple[str, str]] = []
    for sub in ("1_loop_4chunk", "3_multi_ctx_4chunk"):
        d = os.path.join(evals_root, sub)
        if not os.path.isdir(d):
            continue
        for root, _dirs, files in os.walk(d):
            for f in files:
                if f.endswith("_gen_only.mp4"):
                    absp = os.path.join(root, f)
                    rel = os.path.relpath(absp, evals_root)
                    out.append((rel, absp))
    return sorted(out, key=lambda x: x[0])


def load_video_frames(path: str, max_frames: int | None = None) -> np.ndarray:
    """
    Load video as array of frames (RGB, uint8).
    Returns (N, H, W, 3). If max_frames set, stop after that many frames.
    """
    if not HAS_CV2:
        raise RuntimeError("opencv-python is required for video loading (pip install opencv-python)")
    cap = cv2.VideoCapture(path)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame_rgb)
        if max_frames is not None and len(frames) >= max_frames:
            break
    cap.release()
    if not frames:
        return np.zeros((0, 0, 0, 3), dtype=np.uint8)
    return np.stack(frames, axis=0)


def load_video_frames_pil(path: str, max_frames: int | None = None):
    """Load video as list of PIL Images (for CLIP etc.)."""
    from PIL import Image
    arr = load_video_frames(path, max_frames=max_frames)
    if arr.size == 0:
        return []
    return [Image.fromarray(arr[i]) for i in range(arr.shape[0])]