import gradio as gr
from PIL import Image
import os
import base64
import io

def add_logo_base64(image_base64, logo_size_percent=20, margin_percent=2):
    """
    Add logo to image provided as base64 string
    
    Args:
        image_base64: Base64 encoded image string (with or without data URL prefix)
        logo_size_percent: Size of logo as percentage of main image width (default 20%)
        margin_percent: Margin from edges as percentage of main image width (default 2%)
    """
    if not image_base64:
        return None
    
    # Handle base64 input
    if image_base64.startswith('data:image'):
        # Remove data URL prefix
        base64_str = image_base64.split(',')[1]
    else:
        base64_str = image_base64
    
    # Decode base64 to image
    try:
        image_data = base64.b64decode(base64_str)
        main_image = Image.open(io.BytesIO(image_data))
    except Exception as e:
        return f"Error decoding image: {str(e)}"
    
    # Load the logo
    logo_path = "logo.png"
    if not os.path.exists(logo_path):
        # Return original as base64 if logo not found
        return image_base64
    
    # Open logo
    logo = Image.open(logo_path).convert("RGBA")
    
    # Convert main image to RGBA to support transparency
    main_img = main_image.convert("RGBA")
    
    # Calculate logo size (maintain aspect ratio)
    logo_width = int(main_img.width * logo_size_percent / 100)
    logo_height = int(logo.height * (logo_width / logo.width))
    
    # Resize logo
    logo = logo.resize((logo_width, logo_height), Image.Resampling.LANCZOS)
    
    # Calculate position (top-right corner with margin)
    margin = int(main_img.width * margin_percent / 100)
    x_position = main_img.width - logo_width - margin
    y_position = margin
    
    # Create a copy of the main image
    result = main_img.copy()
    
    # Paste logo with transparency
    result.paste(logo, (x_position, y_position), logo)
    
    # Convert back to RGB
    result = result.convert("RGB")
    
    # Convert result to base64
    buffered = io.BytesIO()
    result.save(buffered, format="PNG")
    img_base64 = base64.b64encode(buffered.getvalue()).decode()
    
    return f"data:image/png;base64,{img_base64}"

# Create simple interface for API
iface = gr.Interface(
    fn=add_logo_base64,
    inputs=gr.Textbox(label="Image Base64", placeholder="Paste base64 encoded image here"),
    outputs=gr.Textbox(label="Result Base64"),
    title="Logo Adder API",
    description="Send base64 encoded image to add logo"
)

if __name__ == "__main__":
    iface.launch()