from __future__ import annotations

from pathlib import Path

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


def load(path: str | Path, color: str = "rgb") -> mx.array:
    image_path = Path(path)
    if not image_path.exists():
        raise FileNotFoundError(f"Image file not found: {image_path}")

    if color not in {"rgb", "bgr", "gray", "rgba"}:
        raise ValueError('color must be one of {"rgb", "bgr", "gray", "rgba"}')

    flag = cv2.IMREAD_COLOR
    if color == "gray":
        flag = cv2.IMREAD_GRAYSCALE
    elif color == "rgba":
        flag = cv2.IMREAD_UNCHANGED

    image = cv2.imread(str(image_path), flag)
    if image is None:
        raise ValueError(f"Failed to decode image: {image_path}")

    if color == "rgb":
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    elif color == "rgba":
        if image.ndim == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGBA)
        elif image.shape[2] == 3:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
        else:
            image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)

    if image.ndim == 2:
        image = image[:, :, np.newaxis]

    return mx.array(image)


def load_batch(
    paths: list[str | Path],
    color: str = "rgb",
    size: tuple[int, int] | None = None,
) -> mx.array:
    from ._transforms import resize as _resize

    if not paths:
        raise ValueError("paths must not be empty")

    images = []
    for path in paths:
        image = load(path, color=color)
        if size is not None:
            image = _resize(image, size)
        images.append(np.asarray(image))

    shapes = {image.shape for image in images}
    if len(shapes) != 1:
        raise ValueError(
            f"Images have mismatched shapes {shapes}. "
            "Provide size= to resize all images to a common shape."
        )

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