from __future__ import annotations

from pathlib import Path

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

import mlx_vision as mv


def _write_test_image(path: Path, height: int = 32, width: int = 48) -> None:
    image = np.random.randint(0, 255, (height, width, 3), dtype=np.uint8)
    cv2.imwrite(str(path), image)


def test_load_rgb(tmp_path: Path) -> None:
    image_path = tmp_path / "img.jpg"
    _write_test_image(image_path)
    image = mv.load(image_path)
    assert image.ndim == 3
    assert image.shape[2] == 3
    assert image.dtype == mx.uint8


def test_load_grayscale(tmp_path: Path) -> None:
    image_path = tmp_path / "img.png"
    _write_test_image(image_path)
    image = mv.load(image_path, color="gray")
    assert image.shape[2] == 1


def test_load_batch(tmp_path: Path) -> None:
    paths = []
    for index in range(3):
        image_path = tmp_path / f"img_{index}.png"
        _write_test_image(image_path, height=16, width=16)
        paths.append(image_path)

    batch = mv.load_batch(paths)
    assert batch.shape == (3, 16, 16, 3)
    assert batch.dtype == mx.uint8


def test_load_batch_requires_consistent_shapes(tmp_path: Path) -> None:
    image_a = tmp_path / "a.png"
    image_b = tmp_path / "b.png"
    _write_test_image(image_a, height=16, width=16)
    _write_test_image(image_b, height=20, width=16)

    with pytest.raises(ValueError):
        mv.load_batch([image_a, image_b])
