from __future__ import annotations

import base64
from io import BytesIO

import numpy as np
import torch
from PIL import Image


def _iter_image_tensors(image: torch.Tensor) -> list[torch.Tensor]:
    if image.ndim == 4:
        return [image[index] for index in range(image.shape[0])]
    if image.ndim == 3:
        return [image]
    raise ValueError("Expected image tensor with shape [B, H, W, C] or [H, W, C].")


def _normalize_image_tensor(image: torch.Tensor) -> torch.Tensor:
    if image.ndim == 4:
        image = image[0]
    if image.ndim != 3:
        raise ValueError("Expected image tensor with shape [B, H, W, C] or [H, W, C].")
    return image.clamp(0.0, 1.0).cpu()


def tensor_to_png_bytes(image: torch.Tensor) -> bytes:
    normalized = _normalize_image_tensor(image)
    image_array = (normalized.numpy() * 255).astype(np.uint8)
    pil_image = Image.fromarray(image_array)
    buffer = BytesIO()
    pil_image.save(buffer, format="PNG")
    return buffer.getvalue()


def tensor_to_base64_png(image: torch.Tensor) -> str:
    raw = tensor_to_png_bytes(image)
    encoded = base64.b64encode(raw).decode("utf-8")
    return f"data:image/png;base64,{encoded}"


def tensor_batch_to_base64_pngs(image: torch.Tensor) -> list[str]:
    return [tensor_to_base64_png(item) for item in _iter_image_tensors(image)]


def base64_data_uri_to_bytes(data_uri: str) -> bytes:
    if "," not in data_uri:
        raise ValueError("Expected a valid data URI.")
    return base64.b64decode(data_uri.split(",", 1)[1])


def decode_image_bytes(raw_bytes: bytes) -> torch.Tensor:
    pil_image = Image.open(BytesIO(raw_bytes)).convert("RGB")
    image_array = np.asarray(pil_image).astype(np.float32) / 255.0
    return torch.from_numpy(image_array).unsqueeze(0)


def stack_image_tensors(images: list[torch.Tensor]) -> torch.Tensor:
    if not images:
        raise ValueError("Expected at least one image tensor to stack.")
    normalized = []
    for image in images:
        if image.ndim == 3:
            normalized.append(image.unsqueeze(0))
        else:
            normalized.append(image)
    return torch.cat(normalized, dim=0)