from __future__ import annotations

from pathlib import Path

import cv2
import mlx.core as mx
import numpy as np


def load(
    path: str | Path,
    fps: float | None = None,
    start_time: float | None = None,
    end_time: float | None = None,
    max_frames: int | None = None,
) -> mx.array:
    """Load an RGB video clip into a THWC MLX array."""
    video_path = Path(path)
    if not video_path.exists():
        raise FileNotFoundError(f"Video file not found: {video_path}")

    capture = cv2.VideoCapture(str(video_path))
    if not capture.isOpened():
        raise ValueError(f"Failed to open video file: {video_path}")

    native_fps = capture.get(cv2.CAP_PROP_FPS)
    if fps is not None and fps <= 0:
        raise ValueError("fps must be positive")
    if start_time is not None and start_time < 0:
        raise ValueError("start_time must be non-negative")
    if end_time is not None and start_time is not None and end_time < start_time:
        raise ValueError("end_time must be greater than or equal to start_time")
    if max_frames is not None and max_frames <= 0:
        raise ValueError("max_frames must be positive")

    if start_time:
        capture.set(cv2.CAP_PROP_POS_MSEC, start_time * 1000.0)

    next_sample_time = start_time or 0.0
    sample_period = None if fps is None else 1.0 / fps
    frames: list[np.ndarray] = []

    while True:
        ok, frame = capture.read()
        if not ok:
            break

        frame_index = capture.get(cv2.CAP_PROP_POS_FRAMES) - 1
        if native_fps and native_fps > 0:
            timestamp = frame_index / native_fps
        else:
            timestamp = capture.get(cv2.CAP_PROP_POS_MSEC) / 1000.0

        if end_time is not None and timestamp > end_time:
            break

        if sample_period is not None and timestamp + 1e-9 < next_sample_time:
            continue

        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

        if sample_period is not None:
            next_sample_time += sample_period
            if timestamp > next_sample_time:
                next_sample_time = timestamp + sample_period

        if max_frames is not None and len(frames) >= max_frames:
            break

    capture.release()

    if not frames:
        raise ValueError("No frames were decoded for the requested video range")

    return mx.array(np.stack(frames, axis=0))
