"""Read what someone uploads into a uint8 array: one 2D layer (H, W) or a layer stack (L, H, W).

Accepted: a single image (PNG, JPEG, WebP, BMP, one-page TIFF), a multi-page TIFF of surface layers, a ZIP of layer
images (ordered by the numbers in their names, as villa and VC3D name them: 00.tif, 01.tif, ...), or an .npz with an
array "volume" (Ink Lens's region format). Color is turned to gray, 16-bit is scaled to 8-bit. Nothing is kept.
"""

from __future__ import annotations

import io
import re
import zipfile
from pathlib import Path

import numpy as np

MIN_SIDE, MAX_SIDE = 128, 2048
MAX_LAYERS = 65
MAX_BYTES = 400 * 1024 * 1024
IMAGE_EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff")


class UploadError(ValueError):
    """Shown to the person as it is."""


def to_u8(a):
    """Gray uint8 from a 2D or (H, W, C) array of any integer or float type."""
    a = np.asarray(a)
    if a.ndim == 3 and a.shape[-1] in (3, 4):
        rgb_u8 = a.dtype == np.uint8
        a = a[..., :3].astype(np.float32) @ np.array([0.299, 0.587, 0.114], np.float32)
        if rgb_u8:
            return np.clip(a + 0.5, 0, 255).astype(np.uint8)
    if a.ndim != 2:
        raise UploadError(f"expected a 2D image, got shape {tuple(a.shape)}")
    if a.dtype == np.uint8:
        return a
    if a.dtype == np.uint16:
        return (a // 257).astype(np.uint8)
    is_float = a.dtype.kind == "f"
    a = np.nan_to_num(a.astype(np.float32))
    lo, hi = float(a.min()), float(a.max())
    if is_float and lo >= 0 and hi <= 1.0:
        return np.clip(a * 255 + 0.5, 0, 255).astype(np.uint8)
    if lo >= 0 and hi <= 255:
        return np.clip(a + 0.5, 0, 255).astype(np.uint8)
    if hi <= lo:
        return np.zeros(a.shape, np.uint8)
    return np.clip((a - lo) / (hi - lo) * 255 + 0.5, 0, 255).astype(np.uint8)


def _image(data, name):
    if name.lower().endswith((".tif", ".tiff")):
        import tifffile
        a = tifffile.imread(io.BytesIO(data))
        if a.ndim == 2 or (a.ndim == 3 and a.shape[-1] in (3, 4)):
            return [to_u8(a)]
        if a.ndim in (3, 4):
            return [to_u8(l) for l in a]
        raise UploadError(f"cannot read a TIFF of shape {tuple(a.shape)} as layers")
    from PIL import Image
    with Image.open(io.BytesIO(data)) as im:
        return [to_u8(np.asarray(im.convert("L") if im.mode not in ("L", "I;16", "I", "F") else im))]


def _natural(name):
    return [int(t) if t.isdigit() else t.lower() for t in re.split(r"(\d+)", Path(name).name)]


def read(path):
    """(array, kind, note): kind is "image" for one layer, "stack" for several."""
    path = Path(path)
    size = path.stat().st_size
    if size > MAX_BYTES:
        raise UploadError(f"the file is {size / 2**20:.0f} MB; the limit is {MAX_BYTES // 2**20} MB")
    name = path.name.lower()
    if name.endswith(".npz"):
        with np.load(path, allow_pickle=False) as z:
            if "volume" not in z:
                raise UploadError('an .npz needs an array named "volume" (layers, height, width)')
            layers = [to_u8(v) for v in np.asarray(z["volume"])]
        note = f"{len(layers)} layers from the .npz"
    elif name.endswith(".zip"):
        with zipfile.ZipFile(path) as zf:
            names = sorted((n for n in zf.namelist() if n.lower().endswith(IMAGE_EXT)
                            and not Path(n).name.startswith(".") and "__MACOSX" not in n), key=_natural)
            if not names:
                raise UploadError("the ZIP holds no layer images (.tif, .png or .jpg)")
            if len(names) > MAX_LAYERS:
                raise UploadError(f"the ZIP holds {len(names)} images; send at most {MAX_LAYERS} layers")
            layers = []
            for n in names:
                layers.extend(_image(zf.read(n), n))
        note = f"{len(layers)} layers from the ZIP, in name order ({Path(names[0]).name} first)"
    elif name.endswith(IMAGE_EXT):
        layers = _image(path.read_bytes(), name)
        note = "one image" if len(layers) == 1 else f"{len(layers)} pages of the TIFF as layers"
    else:
        raise UploadError("send an image (PNG, JPEG, TIFF), a multi-page TIFF or ZIP of layers, or an .npz")
    shapes = {l.shape for l in layers}
    if len(shapes) != 1:
        raise UploadError(f"the layers differ in size: {sorted(shapes)}")
    h, w = layers[0].shape
    if min(h, w) < MIN_SIDE:
        raise UploadError(f"the image is {w} x {h} px; it needs at least {MIN_SIDE} px on each side")
    if len(layers) > MAX_LAYERS:
        raise UploadError(f"{len(layers)} layers; send at most {MAX_LAYERS}")
    if max(h, w) > MAX_SIDE:
        raise UploadError(f"the image is {w} x {h} px; crop it to at most {MAX_SIDE} px on each side")
    if len(layers) == 1:
        return layers[0], "image", note
    return np.stack(layers), "stack", note