import cv2
import numpy as np
from PIL import Image
import io
import base64
from typing import Tuple, Optional


def load_image_from_bytes(image_bytes: bytes) -> np.ndarray:
    """
    Load image from bytes into numpy array

    Args:
        image_bytes: Raw image bytes

    Returns:
        Image as RGB numpy array
    """
    image = Image.open(io.BytesIO(image_bytes))

    # Convert to RGB if needed
    if image.mode != 'RGB':
        image = image.convert('RGB')

    return np.array(image)


def load_image_from_base64(base64_string: str) -> np.ndarray:
    """
    Load image from base64 string

    Args:
        base64_string: Base64 encoded image

    Returns:
        Image as RGB numpy array
    """
    # Remove data URL prefix if present
    if ',' in base64_string:
        base64_string = base64_string.split(',')[1]

    image_bytes = base64.b64decode(base64_string)
    return load_image_from_bytes(image_bytes)


def resize_image(
    image: np.ndarray,
    target_size: int,
    maintain_aspect: bool = True
) -> Tuple[np.ndarray, Tuple[int, int]]:
    """
    Resize image to target size

    Args:
        image: Input image array
        target_size: Target size (will be longest edge if maintain_aspect=True)
        maintain_aspect: Whether to maintain aspect ratio

    Returns:
        Tuple of (resized_image, original_size)
    """
    h, w = image.shape[:2]
    original_size = (w, h)

    if maintain_aspect:
        # Calculate new dimensions maintaining aspect ratio
        if h > w:
            new_h = target_size
            new_w = int(w * (target_size / h))
        else:
            new_w = target_size
            new_h = int(h * (target_size / w))
    else:
        new_w = target_size
        new_h = target_size

    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    return resized, original_size


def normalize_image(image: np.ndarray) -> np.ndarray:
    """
    Normalize image for model input

    Args:
        image: Input image array (RGB)

    Returns:
        Normalized image array
    """
    # Convert to float32 and normalize to [0, 1]
    image = image.astype(np.float32) / 255.0

    # ImageNet normalization
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])

    image = (image - mean) / std

    return image


def depth_to_colormap(
    depth: np.ndarray,
    colormap: int = cv2.COLORMAP_INFERNO
) -> np.ndarray:
    """
    Convert depth map to colorized visualization

    Args:
        depth: Depth map array
        colormap: OpenCV colormap constant

    Returns:
        Colorized depth map (RGB)
    """
    # Normalize depth to 0-255
    depth_normalized = cv2.normalize(depth, None, 0, 255, cv2.NORM_MINMAX)
    depth_uint8 = depth_normalized.astype(np.uint8)

    # Apply colormap
    colored = cv2.applyColorMap(depth_uint8, colormap)

    # Convert BGR to RGB
    colored = cv2.cvtColor(colored, cv2.COLOR_BGR2RGB)

    return colored


def array_to_base64(image: np.ndarray, format: str = 'PNG') -> str:
    """
    Convert numpy array to base64 string

    Args:
        image: Image array
        format: Output format (PNG, JPEG, etc.)

    Returns:
        Base64 encoded image string
    """
    pil_image = Image.fromarray(image.astype(np.uint8))

    buffer = io.BytesIO()
    pil_image.save(buffer, format=format)
    buffer.seek(0)

    base64_string = base64.b64encode(buffer.read()).decode('utf-8')
    return f"data:image/{format.lower()};base64,{base64_string}"


def array_to_bytes(image: np.ndarray, format: str = 'PNG') -> bytes:
    """
    Convert numpy array to bytes

    Args:
        image: Image array
        format: Output format (PNG, JPEG, etc.)

    Returns:
        Image bytes
    """
    pil_image = Image.fromarray(image.astype(np.uint8))

    buffer = io.BytesIO()
    pil_image.save(buffer, format=format)
    buffer.seek(0)

    return buffer.read()


def create_side_by_side(
    original: np.ndarray,
    depth: np.ndarray,
    colormap: bool = True
) -> np.ndarray:
    """
    Create side-by-side comparison of original and depth

    Args:
        original: Original image
        depth: Depth map
        colormap: Whether to apply colormap to depth

    Returns:
        Side-by-side image
    """
    # Ensure same height
    h = original.shape[0]
    depth_resized = cv2.resize(depth, (depth.shape[1], h))

    if colormap and len(depth_resized.shape) == 2:
        depth_resized = depth_to_colormap(depth_resized)

    # Concatenate horizontally
    combined = np.hstack([original, depth_resized])

    return combined