from PIL import Image, ImageOps
import io

def parse_aspect_ratio(aspect_ratio_str: str) -> float:
    """
    Parse aspect ratio string like "16:9" or "9:16" to float.
    """
    parts = aspect_ratio_str.split(':')
    if len(parts) == 2:
        return float(parts[0]) / float(parts[1])
    return 16.0 / 9.0  # fallback

def _parse_hex_color(fill_color: str) -> tuple:
    """
    Safely parse hex color like "#RRGGBB" to an RGB tuple for PIL.
    Falls back to black if parsing fails.
    """
    if not fill_color:
        return (0, 0, 0)
    s = fill_color.strip()
    try:
        if s.startswith("#"):
            s = s[1:]
        if len(s) == 6:
            r = int(s[0:2], 16)
            g = int(s[2:4], 16)
            b = int(s[4:6], 16)
            return (r, g, b)
    except Exception:
        pass
    return (0, 0, 0)

def process_image(image_bytes: bytes, mode: str, fill_color: str = "#000000", target_aspect_ratio: str = "16:9") -> bytes:
    """
    Processes an image to fit the specified aspect ratio.

    Args:
        image_bytes: The original image in bytes.
        mode: The processing mode ('stretch_to_fill', 'compress_to_fit', 'fill').
        fill_color: The hex color code for the 'fill' mode background.
        target_aspect_ratio: The target aspect ratio string like "16:9" or "9:16".

    Returns:
        The processed image in bytes (JPEG format).
    """
    if not mode or mode == 'none':
        return image_bytes

    try:
        target_ratio = parse_aspect_ratio(target_aspect_ratio)
        image = Image.open(io.BytesIO(image_bytes))
        # Normalize EXIF orientation to ensure width/height reflect actual display
        try:
            image = ImageOps.exif_transpose(image)
        except Exception:
            pass
        original_width, original_height = image.size
        original_aspect_ratio = original_width / original_height

        # If aspect ratio is already correct, no processing needed
        if abs(original_aspect_ratio - target_ratio) < 1e-6:
            return image_bytes

        if mode == 'stretch_to_fill':
            if original_aspect_ratio > target_ratio: # Wider than target
                # Height is relatively too short, stretch it
                new_width = original_width
                new_height = int(new_width / target_ratio)
                processed_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
            else: # Taller than target
                # Width is relatively too short, stretch it
                new_height = original_height
                new_width = int(new_height * target_ratio)
                processed_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

        elif mode == 'compress_to_fit':
            if original_aspect_ratio > target_ratio: # Wider than target
                # Width is relatively too long, compress it
                new_height = original_height
                new_width = int(new_height * target_ratio)
                processed_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
            else: # Taller than target
                # Height is relatively too long, compress it
                new_width = original_width
                new_height = int(new_width / target_ratio)
                processed_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

        elif mode == 'fill':
            # Pad with color to meet target aspect ratio without scaling the source
            if original_aspect_ratio > target_ratio:
                # Wider than target -> pad top/bottom
                new_width = original_width
                new_height = int(round(new_width / target_ratio))
                paste_x, paste_y = 0, (new_height - original_height) // 2
            else:
                # Taller than target -> pad left/right
                new_height = original_height
                new_width = int(round(new_height * target_ratio))
                paste_x, paste_y = (new_width - original_width) // 2, 0

            # Use parsed RGB to avoid mode issues and ensure consistent color
            fill_rgb = _parse_hex_color(fill_color or "#000000")
            background = Image.new('RGB', (new_width, new_height), fill_rgb)

            # Preserve alpha if present so the background color shows through transparent regions
            if image.mode in ('RGBA', 'LA') or (image.mode == 'P' and 'transparency' in image.info):
                if image.mode != 'RGBA':
                    image = image.convert('RGBA')
                background.paste(image, (paste_x, paste_y), image)
            else:
                background.paste(image, (paste_x, paste_y))
            processed_image = background
        else:
             # If mode is unknown, return original image
            return image_bytes


        # Save the processed image to a byte buffer
        byte_arr = io.BytesIO()
        # Ensure image is in RGB mode before saving as JPEG
        if processed_image.mode != 'RGB':
            processed_image = processed_image.convert('RGB')
        processed_image.save(byte_arr, format='JPEG')
        return byte_arr.getvalue()

    except Exception as e:
        print(f"Error processing image: {e}")
        # In case of error, return the original image bytes
        return image_bytes