"""Image and tensor conversion helpers shared by LayerForge endpoints."""

import numpy as np
from PIL import Image
from torchvision import transforms

from .image_serialization import data_url_to_pil, pil_to_data_url
from .node import log


def convert_base64_to_tensor(base64_str):
    """Convert a data-URL image into a BCHW tensor and optional alpha tensor."""
    try:
        img = data_url_to_pil(base64_str)

        has_alpha = img.mode == "RGBA"
        alpha = None
        if has_alpha:
            alpha = img.split()[3]
            background = Image.new("RGB", img.size, (255, 255, 255))
            background.paste(img, mask=alpha)
            img = background
        elif img.mode != "RGB":
            img = img.convert("RGB")

        img_tensor = transforms.ToTensor()(img).unsqueeze(0)
        if has_alpha:
            alpha_tensor = transforms.ToTensor()(alpha).unsqueeze(0)
            return img_tensor, alpha_tensor

        return img_tensor, None
    except Exception as error:
        log.error(f"Error in convert_base64_to_tensor: {error}")
        raise


def convert_tensor_to_base64(tensor, alpha_mask=None, original_alpha=None):
    """Convert a tensor to a PNG data URL, optionally preserving alpha."""
    try:
        tensor = tensor.cpu()

        if tensor.dim() == 4:
            tensor = tensor.squeeze(0)
        if tensor.dim() == 3 and tensor.shape[0] in [1, 3]:
            tensor = tensor.permute(1, 2, 0)

        img_array = (tensor.numpy() * 255).astype(np.uint8)

        if alpha_mask is not None and original_alpha is not None:
            alpha_mask = (alpha_mask.cpu().squeeze().numpy() * 255).astype(np.uint8)
            original_alpha = (original_alpha.cpu().squeeze().numpy() * 255).astype(np.uint8)
            combined_alpha = np.minimum(alpha_mask, original_alpha)

            img = Image.fromarray(img_array, mode="RGB")
            img.putalpha(Image.fromarray(combined_alpha, mode="L"))
        elif img_array.shape[-1] == 1:
            img = Image.fromarray(img_array.squeeze(-1), mode="L")
        else:
            img = Image.fromarray(img_array, mode="RGB")

        return pil_to_data_url(img)
    except Exception as error:
        log.error(f"Error in convert_tensor_to_base64: {error}")
        log.debug(f"Tensor shape: {tensor.shape}, dtype: {tensor.dtype}")
        raise


__all__ = ["convert_base64_to_tensor", "convert_tensor_to_base64"]