import io
import os
import zipfile
from typing import List, Tuple, Union

import numpy as np
from PIL import Image

# ----------------------------
# I/O helpers
# ----------------------------
BytesLike = Union[bytes, bytearray, io.BytesIO]


def load_image(src: Union[str, BytesLike]) -> Image.Image:
    """Load image from path or in-memory bytes → RGB PIL.Image."""
    if isinstance(src, (bytes, bytearray, io.BytesIO)):
        return Image.open(io.BytesIO(src)).convert("RGB")
    return Image.open(src).convert("RGB")


def unzip_images(file_bytes: bytes, exts=(".png", ".jpg", ".jpeg", ".tif", ".tiff")) -> List[Tuple[str, Image.Image]]:
    """Return [(name, PIL RGB)] from a ZIP bytes blob, filtering by extension."""
    out = []
    with zipfile.ZipFile(io.BytesIO(file_bytes)) as zf:
        for name in zf.namelist():
            if name.lower().endswith(exts) and not name.endswith("/"):
                with zf.open(name) as fh:
                    img = Image.open(io.BytesIO(fh.read())).convert("RGB")
                    out.append((name, img))
    return out


def ensure_dir(path: str) -> None:
    os.makedirs(path, exist_ok=True)


# ----------------------------
# Array utilities
# ----------------------------
def pil_to_np(img: Image.Image) -> np.ndarray:
    """PIL RGB → np.uint8 [H, W, 3]."""
    return np.asarray(img, dtype=np.uint8)


def np_to_pil(arr: np.ndarray) -> Image.Image:
    """np.uint8 [H, W, 3] → PIL RGB."""
    if arr.dtype != np.uint8:
        arr = np.clip(arr, 0, 255).astype(np.uint8)
    return Image.fromarray(arr, mode="RGB")


def normalize01(arr: np.ndarray) -> np.ndarray:
    """Scale array to [0,1] with safe denominator."""
    arr = arr.astype(np.float32)
    mn = float(arr.min())
    mx = float(arr.max())
    denom = (mx - mn) if (mx - mn) != 0 else 1.0
    return (arr - mn) / denom


# ----------------------------
# Heatmap coloring
# ----------------------------
def colorize_heatmap(attn_grid: np.ndarray) -> Image.Image:
    """
    Convert [H,W] float in [0,1] to a colored heatmap RGB PIL image using JET.
    """
    import cv2

    attn = np.clip(attn_grid, 0.0, 1.0)
    grid8 = (attn * 255).astype(np.uint8)
    cm_bgr = cv2.applyColorMap(grid8, cv2.COLORMAP_JET)
    cm_rgb = cv2.cvtColor(cm_bgr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(cm_rgb, mode="RGB")