from __future__ import annotations

import functools
from typing import (
    Any,
    cast,
)

import cv2
import expath
import numpy as np
import PIL.Image as pil_image
import torch
import torchvision.tv_tensors

__all__ = ["get_kwd", "read_pixels", "write_png_l16", "write_png_rgb"]

transform_context = functools.partial(
    torchvision.tv_tensors.set_return_type, "TVTensor"
)


def read_pixels(path: str, color: bool, alpha=False) -> torch.Tensor:
    """Read an image file using OpenCV."""
    if alpha:
        flags = cv2.IMREAD_UNCHANGED
    else:
        flags = (
            cv2.IMREAD_COLOR if color else cv2.IMREAD_GRAYSCALE
        ) | cv2.IMREAD_ANYDEPTH

    image = cv2.imread(path, flags)
    if image is None:
        msg = f"Failed to read image at path: {path}"
        raise RuntimeError(msg)

    if color and alpha:
        image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
    elif color:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    return torch.from_numpy(image.astype(np.int32))


def write_png_rgb(path: expath.PathType, tensor: torch.Tensor) -> None:
    """Write a tensor to a PNG file (RGB)."""
    assert tensor.ndim == 3
    if tensor.shape[0] == 3:
        tensor = tensor.permute(1, 2, 0)  # CHW -> HWC
    elif tensor.shape[-1] != 3:
        msg = "Input tensor must have 3 channels (CHW or HWC)"
        raise ValueError(msg)

    path = expath.locate(path)
    image = tensor.cpu().numpy().astype("uint8")
    pil_image.fromarray(image, mode="RGB").save(path)


def write_png_l16(path: expath.PathType, tensor: torch.Tensor) -> None:
    """Write a tensor to a PNG file (L; 16-bit)."""
    path = expath.locate(path)
    image = tensor.cpu().numpy().astype("uint16")
    pil_image.fromarray(image).save(path)


def get_kwd[KeywordType](
    kwds_dict: dict[str, Any], name: str, cast_as: type[KeywordType] | None = None, /
) -> KeywordType:
    try:
        res = kwds_dict.pop(name)
    except KeyError as e:
        msg = f"Missing keyword argument: {e.args[0]!r}"
        raise TypeError(msg) from e

    return cast(KeywordType, res)
