"""
Image compositing and enhancement module.
"""

from PIL import Image, ImageEnhance, ImageFilter
import numpy as np


def composite_image(foreground, background, target_size=(1200, 1600)):
    """
    Composite the animal (with removed background) onto a new background.
    
    Args:
        foreground: PIL Image (RGBA) of animal with transparent background
        background: PIL Image of the background scene
        target_size: tuple of (width, height) for final output
        
    Returns:
        PIL Image: Composited and enhanced image
    """
    
    # Resize background to target size
    background = background.convert('RGB')
    background = background.resize(target_size, Image.Resampling.LANCZOS)
    
    # Calculate foreground size (keep animal prominent but not too large)
    # Animal should take up about 60-70% of the height
    fg_max_height = int(target_size[1] * 0.7)
    fg_max_width = int(target_size[0] * 0.8)
    
    # Resize foreground maintaining aspect ratio
    fg_ratio = foreground.size[0] / foreground.size[1]
    
    if foreground.size[1] > fg_max_height:
        new_height = fg_max_height
        new_width = int(new_height * fg_ratio)
    else:
        new_height = foreground.size[1]
        new_width = foreground.size[0]
    
    # Ensure width doesn't exceed max
    if new_width > fg_max_width:
        new_width = fg_max_width
        new_height = int(new_width / fg_ratio)
    
    foreground_resized = foreground.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # Calculate position to place the animal at the bottom (grounded, no margin)
    x_pos = (target_size[0] - new_width) // 2  # Center horizontally
    y_pos = target_size[1] - new_height  # Bottom edge, no margin
    
    # Ensure position is within bounds
    y_pos = max(0, min(y_pos, target_size[1] - new_height))
    
    # Create a copy of background to paste onto
    result = background.copy()
    
    # Paste foreground using alpha channel as mask
    if foreground_resized.mode == 'RGBA':
        result.paste(foreground_resized, (x_pos, y_pos), foreground_resized)
    else:
        # If no alpha channel, just paste normally
        result.paste(foreground_resized, (x_pos, y_pos))
    
    # Apply enhancements
    result = enhance_image(result)
    
    return result


def enhance_image(image):
    """
    Apply subtle enhancements to make the image more appealing.
    
    Args:
        image: PIL Image (RGB)
        
    Returns:
        PIL Image: Enhanced image
    """
    
    # Increase brightness slightly (5%)
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(1.05)
    
    # Increase contrast slightly (10%)
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(1.10)
    
    # Increase color saturation slightly (8%)
    enhancer = ImageEnhance.Color(image)
    image = enhancer.enhance(1.08)
    
    # Apply subtle sharpening
    image = image.filter(ImageFilter.UnsharpMask(radius=1, percent=50, threshold=3))
    
    return image


def blend_edges(foreground, background, blur_radius=5):
    """
    Optional: Blend the edges of the foreground with the background for smoother compositing.
    
    Args:
        foreground: PIL Image (RGBA)
        background: PIL Image (RGB)
        blur_radius: int, radius for edge blurring
        
    Returns:
        PIL Image: Composited image with blended edges
    """
    
    if foreground.mode != 'RGBA':
        return foreground
    
    # Extract alpha channel
    alpha = foreground.split()[3]
    
    # Apply slight blur to alpha for softer edges
    alpha_blurred = alpha.filter(ImageFilter.GaussianBlur(blur_radius))
    
    # Recreate foreground with blurred alpha
    fg_with_blur = Image.merge('RGBA', foreground.split()[:3] + (alpha_blurred,))
    
    return fg_with_blur