image_processor.py

# image_processor.py
"""
Image processing module for FootFa Telegram News Bot
Handles branding overlay application on news post images
"""

import os
import io
import base64
import aiohttp
import asyncio
from PIL import Image, ImageDraw
from typing import Optional

async def apply_branding_overlay(base_image_url: str, overlay_image_path: str) -> Optional[str]:
    """
    Apply a transparent branding overlay to a base image.
    
    Args:
        base_image_url (str): URL or base64 data URI of the base image
        overlay_image_path (str): Path to the overlay PNG file
        
    Returns:
        Optional[str]: Base64 data URI of the composited image, or None if failed
    """
    try:
        print(f"🎨 Applying branding overlay to image...")
        
        # Load base image
        base_image = await load_image_from_url_or_base64(base_image_url)
        if not base_image:
            print("❌ Failed to load base image")
            return None
        
        # Load overlay image
        overlay_image = load_overlay_image(overlay_image_path)
        if not overlay_image:
            print("❌ Failed to load overlay image")
            return None
        
        # Ensure base image is in RGBA mode for transparency handling
        if base_image.mode != 'RGBA':
            base_image = base_image.convert('RGBA')
        
        # Resize overlay to match base image width while maintaining aspect ratio
        overlay_resized = resize_overlay_to_fit(overlay_image, base_image.size)
        
        # Create final composited image
        composited_image = composite_images(base_image, overlay_resized)
        
        # Convert to base64 data URI
        result_data_uri = image_to_base64_data_uri(composited_image)
        
        print(f"✅ Branding overlay applied successfully")
        return result_data_uri
        
    except Exception as e:
        print(f"❌ Error applying branding overlay: {e}")
        return None

async def load_image_from_url_or_base64(image_input: str) -> Optional[Image.Image]:
    """
    Load an image from either a URL or base64 data URI.
    
    Args:
        image_input (str): URL or base64 data URI
        
    Returns:
        Optional[Image.Image]: Loaded PIL Image or None if failed
    """
    try:
        if image_input.startswith('data:image/'):
            # Handle base64 data URI
            print("📥 Loading image from base64 data URI...")
            header, encoded = image_input.split(',', 1)
            image_data = base64.b64decode(encoded)
            return Image.open(io.BytesIO(image_data))
        
        elif image_input.startswith(('http://', 'https://')):
            # Handle HTTP/HTTPS URL with SSL verification disabled
            print(f"📥 Loading image from URL: {image_input[:50]}...")
            
            # Configure connector with SSL verification disabled
            connector = aiohttp.TCPConnector(ssl=False)
            timeout = aiohttp.ClientTimeout(total=30)
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
                'Accept': 'image/webp,image/apng,image/*,*/*;q=0.8',
                'Accept-Language': 'en-US,en;q=0.9',
                'Accept-Encoding': 'gzip, deflate, br',
                'Connection': 'keep-alive',
                'Upgrade-Insecure-Requests': '1',
            }
            
            async with aiohttp.ClientSession(connector=connector, timeout=timeout, headers=headers) as session:
                async with session.get(image_input) as response:
                    if response.status == 200:
                        image_data = await response.read()
                        return Image.open(io.BytesIO(image_data))
                    else:
                        print(f"❌ HTTP {response.status} when fetching image")
                        return None
        else:
            print(f"❌ Unsupported image input format: {image_input[:50]}...")
            return None
            
    except Exception as e:
        print(f"❌ Error loading image: {e}")
        return None

def load_overlay_image(overlay_path: str) -> Optional[Image.Image]:
    """
    Load the overlay image from local file system.
    
    Args:
        overlay_path (str): Path to the overlay PNG file
        
    Returns:
        Optional[Image.Image]: Loaded PIL Image or None if failed
    """
    try:
        if not os.path.exists(overlay_path):
            print(f"❌ Overlay image not found: {overlay_path}")
            return None
        
        print(f"📥 Loading overlay image: {overlay_path}")
        overlay_image = Image.open(overlay_path)
        
        # Ensure overlay is in RGBA mode for transparency
        if overlay_image.mode != 'RGBA':
            overlay_image = overlay_image.convert('RGBA')
        
        print(f"✅ Overlay loaded: {overlay_image.size[0]}x{overlay_image.size[1]} pixels")
        return overlay_image
        
    except Exception as e:
        print(f"❌ Error loading overlay image: {e}")
        return None

def resize_overlay_to_fit(overlay_image: Image.Image, base_size: tuple) -> Image.Image:
    """
    Resize overlay image to fit the base image dimensions.
    Overlay is scaled to match base image width while maintaining aspect ratio.
    
    Args:
        overlay_image (Image.Image): The overlay image to resize
        base_size (tuple): (width, height) of the base image
        
    Returns:
        Image.Image: Resized overlay image
    """
    try:
        base_width, base_height = base_size
        overlay_width, overlay_height = overlay_image.size
        
        # Calculate new dimensions to match base image width
        scale_factor = base_width / overlay_width
        new_overlay_width = base_width
        new_overlay_height = int(overlay_height * scale_factor)
        
        # Resize overlay maintaining aspect ratio
        resized_overlay = overlay_image.resize(
            (new_overlay_width, new_overlay_height), 
            Image.LANCZOS
        )
        
        print(f"🔄 Overlay resized from {overlay_width}x{overlay_height} to {new_overlay_width}x{new_overlay_height}")
        return resized_overlay
        
    except Exception as e:
        print(f"❌ Error resizing overlay: {e}")
        return overlay_image

def composite_images(base_image: Image.Image, overlay_image: Image.Image) -> Image.Image:
    """
    Composite overlay image onto base image with proper transparency handling.
    Overlay is positioned at the bottom of the base image.
    
    Args:
        base_image (Image.Image): The base image
        overlay_image (Image.Image): The overlay image to composite
        
    Returns:
        Image.Image: Composited image
    """
    try:
        base_width, base_height = base_image.size
        overlay_width, overlay_height = overlay_image.size
        
        # Calculate position to place overlay at the bottom
        x_position = 0  # Align to left (overlay should span full width)
        y_position = base_height - overlay_height  # Position at bottom
        
        # If overlay is taller than base image, crop it
        if overlay_height > base_height:
            print("⚠️  Overlay is taller than base image, cropping to fit")
            crop_top = overlay_height - base_height
            overlay_image = overlay_image.crop((0, crop_top, overlay_width, overlay_height))
            y_position = 0
        
        # Create a copy of base image to avoid modifying original
        result_image = base_image.copy()
        
        # Composite overlay onto base image with alpha blending
        if overlay_image.mode == 'RGBA':
            # Use alpha_composite for proper transparency handling
            overlay_positioned = Image.new('RGBA', base_image.size, (0, 0, 0, 0))
            overlay_positioned.paste(overlay_image, (x_position, y_position))
            result_image = Image.alpha_composite(result_image, overlay_positioned)
        else:
            # Fallback to paste method
            result_image.paste(overlay_image, (x_position, y_position), overlay_image)
        
        print(f"✅ Images composited successfully. Overlay positioned at ({x_position}, {y_position})")
        return result_image
        
    except Exception as e:
        print(f"❌ Error compositing images: {e}")
        return base_image

def image_to_base64_data_uri(image: Image.Image, format: str = 'PNG') -> str:
    """
    Convert PIL Image to base64 data URI string.
    
    Args:
        image (Image.Image): PIL Image to convert
        format (str): Image format (PNG, JPEG, etc.)
        
    Returns:
        str: Base64 data URI string
    """
    try:
        # Convert to RGB if saving as JPEG (JPEG doesn't support transparency)
        if format.upper() == 'JPEG' and image.mode == 'RGBA':
            # Create white background for JPEG
            background = Image.new('RGB', image.size, (255, 255, 255))
            background.paste(image, mask=image.split()[-1])  # Use alpha channel as mask
            image = background
        
        # Save image to bytes buffer
        buffer = io.BytesIO()
        image.save(buffer, format=format, quality=95, optimize=True)
        buffer.seek(0)
        
        # Encode to base64
        image_bytes = buffer.getvalue()
        base64_string = base64.b64encode(image_bytes).decode('utf-8')
        
        # Create data URI
        mime_type = f"image/{format.lower()}"
        data_uri = f"data:{mime_type};base64,{base64_string}"
        
        print(f"✅ Image converted to base64 data URI ({len(data_uri)} characters)")
        return data_uri
        
    except Exception as e:
        print(f"❌ Error converting image to base64: {e}")
        # Return a placeholder data URI in case of error
        return "data:image/png;base64,"

# --- Testing Function ---
async def test_branding_overlay():
    """
    Test function to verify branding overlay functionality.
    """
    print("\n=== Testing Branding Overlay Functionality ===")
    
    # Test with a sample image URL
    test_image_url = "https://via.placeholder.com/800x600/FF5722/white?text=FootFa+News"
    overlay_path = "asset/cover.png"
    
    print(f"🧪 Testing with image: {test_image_url}")
    print(f"🧪 Testing with overlay: {overlay_path}")
    
    result = await apply_branding_overlay(test_image_url, overlay_path)
    
    if result:
        print("✅ Branding overlay test successful!")
        print(f"📊 Result data URI length: {len(result)} characters")
    else:
        print("❌ Branding overlay test failed!")

if __name__ == "__main__":
    # Run test when script is executed directly
    asyncio.run(test_branding_overlay())