import torch
import numpy as np
import base64
import io
from PIL import Image
from typing import Union, Tuple

def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
    """Convert ComfyUI tensor to PIL Image"""
    # ComfyUI images are [B, H, W, C] or [H, W, C] tensors
    if len(tensor.shape) == 4:
        # Take first image from batch
        np_image = tensor[0].cpu().numpy()
    else:
        np_image = tensor.cpu().numpy()

    # Convert from float [0,1] to uint8 [0,255]
    if np_image.dtype == np.float32 or np_image.dtype == np.float64:
        np_image = (np_image * 255).astype(np.uint8)

    # Ensure we have 3 channels (RGB)
    if np_image.shape[-1] == 1:
        np_image = np.repeat(np_image, 3, axis=-1)
    elif np_image.shape[-1] == 4:
        # Convert RGBA to RGB
        np_image = np_image[:, :, :3]

    return Image.fromarray(np_image)

def pil_to_tensor(pil_image: Image.Image) -> torch.Tensor:
    """Convert PIL Image to ComfyUI tensor format"""
    # Ensure RGB format
    if pil_image.mode != "RGB":
        pil_image = pil_image.convert("RGB")

    # Convert to numpy array
    np_image = np.array(pil_image).astype(np.float32) / 255.0

    # Convert to tensor with batch dimension: [1, H, W, C]
    tensor = torch.from_numpy(np_image).unsqueeze(0)

    return tensor

def image_to_base64(image: Union[torch.Tensor, Image.Image], format: str = "PNG") -> str:
    """Convert image to base64 string"""
    if isinstance(image, torch.Tensor):
        pil_image = tensor_to_pil(image)
    else:
        pil_image = image

    buffer = io.BytesIO()
    pil_image.save(buffer, format=format)
    image_bytes = buffer.getvalue()

    return base64.b64encode(image_bytes).decode('utf-8')

def base64_to_image(b64_string: str) -> torch.Tensor:
    """Convert base64 string to ComfyUI tensor"""
    try:
        image_bytes = base64.b64decode(b64_string)
        pil_image = Image.open(io.BytesIO(image_bytes))
        return pil_to_tensor(pil_image)
    except Exception as e:
        raise ValueError(f"Failed to decode base64 image: {e}")

def resize_image_aspect_ratio(image: Image.Image, target_ratio: str, max_size: int = 1024) -> Image.Image:
    """Resize image to target aspect ratio while maintaining quality"""
    # Parse aspect ratio
    if target_ratio == "1:1":
        target_width, target_height = 1, 1
    elif target_ratio == "16:9":
        target_width, target_height = 16, 9
    elif target_ratio == "9:16":
        target_width, target_height = 9, 16
    elif target_ratio == "4:3":
        target_width, target_height = 4, 3
    elif target_ratio == "3:4":
        target_width, target_height = 3, 4
    else:
        return image  # Return original if unknown ratio

    # Calculate target dimensions
    target_aspect = target_width / target_height
    current_width, current_height = image.size
    current_aspect = current_width / current_height

    if current_aspect > target_aspect:
        # Image is wider than target, crop width
        new_width = int(current_height * target_aspect)
        left = (current_width - new_width) // 2
        image = image.crop((left, 0, left + new_width, current_height))
    elif current_aspect < target_aspect:
        # Image is taller than target, crop height
        new_height = int(current_width / target_aspect)
        top = (current_height - new_height) // 2
        image = image.crop((0, top, current_width, top + new_height))

    # Resize to max_size while maintaining aspect ratio
    current_width, current_height = image.size
    if max(current_width, current_height) > max_size:
        if current_width > current_height:
            new_width = max_size
            new_height = int(max_size / target_aspect)
        else:
            new_height = max_size
            new_width = int(max_size * target_aspect)

        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

    return image