from __future__ import annotations

from pathlib import Path

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

import mlx_video as mv


def _write_test_video(path: Path, fps: int = 4, frames: int = 8) -> None:
    size = (24, 16)
    writer = cv2.VideoWriter(
        str(path),
        cv2.VideoWriter_fourcc(*"mp4v"),
        fps,
        size,
    )
    if not writer.isOpened():
        raise RuntimeError("Failed to create test video")

    for index in range(frames):
        frame = np.zeros((size[1], size[0], 3), dtype=np.uint8)
        frame[..., 0] = index * 20
        frame[..., 1] = 50
        frame[..., 2] = 200 - index * 10
        writer.write(frame)

    writer.release()


def test_load_returns_rgb_uint8_clip(tmp_path: Path) -> None:
    video_path = tmp_path / "clip.mp4"
    _write_test_video(video_path)

    frames = mv.load(video_path)

    assert frames.shape == (8, 16, 24, 3)
    assert frames.dtype == mx.uint8


def test_load_supports_fps_sampling_and_max_frames(tmp_path: Path) -> None:
    video_path = tmp_path / "clip.mp4"
    _write_test_video(video_path, fps=4, frames=8)

    frames = mv.load(video_path, fps=2, max_frames=3)

    assert frames.shape[0] == 3


def test_resize_to_float_and_normalize() -> None:
    frames = mx.array(np.full((2, 8, 10, 3), 255, dtype=np.uint8))

    resized = mv.resize(frames, (4, 6))
    floats = mv.to_float(resized)
    normalized = mv.normalize(resized)

    assert resized.shape == (2, 4, 6, 3)
    assert floats.dtype == mx.float32
    assert float(mx.min(floats).item()) == pytest.approx(1.0)
    assert normalized.dtype == mx.float32
    assert normalized.shape == (2, 4, 6, 3)


def test_invalid_inputs_raise(tmp_path: Path) -> None:
    with pytest.raises(FileNotFoundError):
        mv.load(tmp_path / "missing.mp4")

    with pytest.raises(ValueError):
        mv.resize(mx.array(np.zeros((4, 4, 3), dtype=np.uint8)), (8, 8))
