from __future__ import annotations

import gzip
import io
import tarfile
from typing import Iterable, List, Tuple

import cv2
import numpy as np


class ImageDecodeError(ValueError):
    """Raised when raw bytes cannot be decoded into a valid image."""
    pass


def maybe_decompress(data: bytes, content_encoding: str | None) -> bytes:
    """Decompress request bytes when `Content-Encoding` indicates gzip."""
    if not content_encoding:
        return data
    encoding = content_encoding.lower().strip()
    if encoding == "gzip":
        # gzip.decompress returns the original raw image/tar bytes.
        return gzip.decompress(data)
    return data


def decode_image_bytes(data: bytes) -> np.ndarray:
    """Decode encoded image bytes into an RGB `numpy.ndarray`."""
    # Build a uint8 view over bytes for OpenCV decoding.
    arr = np.frombuffer(data, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise ImageDecodeError("Failed to decode image bytes")
    # OpenCV decodes BGR by default; convert to RGB for model preprocessing.
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img


def load_webdataset_images(data: bytes) -> List[np.ndarray]:
    """Read a TAR archive and decode every file that looks like an image."""
    images: List[np.ndarray] = []
    with tarfile.open(fileobj=io.BytesIO(data), mode="r:*") as tar:
        for member in tar.getmembers():
            if not member.isfile():
                continue
            fileobj = tar.extractfile(member)
            if fileobj is None:
                continue
            content = fileobj.read()
            try:
                images.append(decode_image_bytes(content))
            # Skip non-image files in the TAR instead of failing all inputs.
            except ImageDecodeError:
                continue
    return images


def resize_and_normalize(
    img: np.ndarray,
    input_size: int,
    mean: Iterable[float],
    std: Iterable[float],
) -> np.ndarray:
    """Resize an RGB image and normalize it into CHW float tensor format."""
    # Resize each image to the fixed square resolution expected by the model.
    resized = cv2.resize(img, (input_size, input_size), interpolation=cv2.INTER_LINEAR)
    # Convert to [0,1] float values before channel-wise normalization.
    tensor = resized.astype(np.float32) / 255.0
    mean_arr = np.array(mean, dtype=np.float32)
    std_arr = np.array(std, dtype=np.float32)
    tensor = (tensor - mean_arr) / std_arr
    # Rearrange from height-width-channel to channel-height-width.
    tensor = np.transpose(tensor, (2, 0, 1))
    return tensor


def prepare_batch(
    images: List[np.ndarray],
    input_size: int,
    mean: Iterable[float],
    std: Iterable[float],
) -> np.ndarray:
    """Convert a list of RGB images into one batched float32 tensor."""
    # Preprocess each image independently, then stack into shape [B, C, H, W].
    batch = [resize_and_normalize(img, input_size, mean, std) for img in images]
    return np.stack(batch, axis=0).astype(np.float32)


def normalize_only(
    img: np.ndarray,
    mean: Iterable[float],
    std: Iterable[float],
) -> np.ndarray:
    """Normalize an RGB image without resizing; returns CHW float32."""
    tensor = img.astype(np.float32) / 255.0
    mean_arr = np.array(mean, dtype=np.float32)
    std_arr = np.array(std, dtype=np.float32)
    tensor = (tensor - mean_arr) / std_arr
    return np.transpose(tensor, (2, 0, 1)).astype(np.float32, copy=False)


def maybe_downscale_max_side(img: np.ndarray, max_side: int) -> tuple[np.ndarray, tuple[int, int]]:
    """Downscale an RGB image so max(H, W) <= max_side, preserving aspect ratio.

    Returns (resized_img, (new_h, new_w)). If no resize is needed, returns the input image.
    """
    if max_side <= 0:
        return img, (int(img.shape[0]), int(img.shape[1]))

    h, w = int(img.shape[0]), int(img.shape[1])
    long_side = max(h, w)
    if long_side <= max_side:
        return img, (h, w)

    scale = float(max_side) / float(long_side)
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))
    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
    return resized, (new_h, new_w)


def prepare_batch_dynamic(
    images: List[np.ndarray],
    mean: Iterable[float],
    std: Iterable[float],
    pad_to: int = 1,
) -> Tuple[np.ndarray, List[Tuple[int, int]]]:
    """Prepare a batch without forcing a fixed resize.

    If images have different H/W, they are padded (bottom/right) to the max H/W in the batch.
    Returns (batch, original_sizes) where sizes are (h, w) in pixels.
    """
    if not images:
        raise ValueError("images must be non-empty")

    orig_sizes: List[Tuple[int, int]] = [(int(im.shape[0]), int(im.shape[1])) for im in images]
    chw = [normalize_only(im, mean, std) for im in images]

    max_h = max(x.shape[1] for x in chw)
    max_w = max(x.shape[2] for x in chw)
    if pad_to and pad_to > 1:
        max_h = ((int(max_h) + int(pad_to) - 1) // int(pad_to)) * int(pad_to)
        max_w = ((int(max_w) + int(pad_to) - 1) // int(pad_to)) * int(pad_to)

    batch = np.zeros((len(chw), 3, max_h, max_w), dtype=np.float32)
    for i, x in enumerate(chw):
        h, w = x.shape[1], x.shape[2]
        batch[i, :, :h, :w] = x

    return batch, orig_sizes