#!/usr/bin/env python3

import os
import google.generativeai as genai
from PIL import Image, ImageDraw, ImageFont
import io
import base64

# Configure API
genai.configure(api_key="AIzaSyCXR9ftl7Rj3i8oNenEBfFQd5JnHdVSI2w")

def resize_image(image_path, target_size=512):
    """Resize image to target size while maintaining aspect ratio."""
    with Image.open(image_path) as img:
        img.thumbnail((target_size, target_size), Image.Resampling.LANCZOS)
        return img.copy()

def load_product_images():
    """Load and resize all Good Bald product images."""
    product_dir = "/Users/igloo/.openclaw/workspace/goodbald-images/original/"
    product_files = [
        "after-shave-treatment-spray.png",
        "post-shave-moisturizer.png", 
        "pre-shave-scrub.png",
        "shave-cream-bentonite-clay.png",
        "shave-cream-tea-tree.png"
    ]
    
    images = []
    for filename in product_files:
        filepath = os.path.join(product_dir, filename)
        resized_img = resize_image(filepath)
        images.append(resized_img)
    
    return images

def generate_base_image():
    """Generate the clean split-screen image using Gemini."""
    
    # Load product images as reference
    product_images = load_product_images()
    
    # Text prompt for Gemini (no text overlays)
    prompt = """Create a high-quality 1080x1080 square split-screen comparison image for a men's grooming brand.

LEFT SIDE (50% of image):
- A 35-40 year old man with patchy, receding hairline and thin hair
- Tired, older appearance with dull, flat lighting
- Slightly desaturated colors
- He looks worn down and less confident

RIGHT SIDE (50% of image): 
- The SAME man but with a freshly shaved bald head
- Clean, confident, healthy appearance
- Warm, golden lighting that makes his skin glow
- He looks younger, more energetic, and confident
- Some of the Good Bald grooming products should be subtly visible in the background/foreground

TECHNICAL REQUIREMENTS:
- Clear visual split down the exact middle
- High quality, professional photography style
- Both sides should clearly be the same person
- Make the transformation dramatic and compelling
- 1080x1080 square format
- NO TEXT OVERLAYS (text will be added separately)

The image should sell the transformation from tired/balding to confident/bald using quality men's grooming products."""

    # Generate image using Gemini
    model = genai.GenerativeModel('gemini-3.1-flash-image-preview')
    
    print("Generating base image with Gemini...")
    
    # Create the content list with images and prompt
    content = product_images + [prompt]
    
    response = model.generate_content(content)
    
    # Save the generated image
    if response.candidates and len(response.candidates) > 0:
        candidate = response.candidates[0]
        if candidate.content and candidate.content.parts:
            for part in candidate.content.parts:
                if hasattr(part, 'inline_data') and part.inline_data:
                    image_data = part.inline_data.data
                    # Convert base64 to image
                    img_bytes = base64.b64decode(image_data)
                    base_img = Image.open(io.BytesIO(img_bytes))
                    return base_img
    
    raise Exception("No image generated from Gemini")

def add_text_overlays(base_image):
    """Add text overlays to the base image using PIL."""
    
    # Create a copy to work with
    img = base_image.copy()
    draw = ImageDraw.Draw(img)
    
    # Try to load a good font, fallback to default
    try:
        # Try system fonts
        font_large = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 48)
        font_small = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 24)
        font_brand = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 28)
    except:
        try:
            font_large = ImageFont.truetype("Arial.ttf", 48)
            font_small = ImageFont.truetype("Arial.ttf", 24)
            font_brand = ImageFont.truetype("Arial.ttf", 28)
        except:
            font_large = ImageFont.load_default()
            font_small = ImageFont.load_default()
            font_brand = ImageFont.load_default()
    
    # Image dimensions (should be 1080x1080)
    width, height = img.size
    
    # Left text: "This adds 5 years" (positioned on left half)
    left_text = "This adds 5 years"
    left_bbox = draw.textbbox((0, 0), left_text, font=font_large)
    left_width = left_bbox[2] - left_bbox[0]
    left_height = left_bbox[3] - left_bbox[1]
    
    left_x = (width // 4) - (left_width // 2)  # Center in left half
    left_y = height - 150  # Near bottom
    
    # Draw text with outline for readability
    # Outline
    for dx, dy in [(-2, -2), (-2, 2), (2, -2), (2, 2)]:
        draw.text((left_x + dx, left_y + dy), left_text, font=font_large, fill="black")
    # Main text
    draw.text((left_x, left_y), left_text, font=font_large, fill="white")
    
    # Right text: "This takes them off" (positioned on right half)
    right_text = "This takes them off"
    right_bbox = draw.textbbox((0, 0), right_text, font=font_large)
    right_width = right_bbox[2] - right_bbox[0]
    right_height = right_bbox[3] - right_bbox[1]
    
    right_x = (3 * width // 4) - (right_width // 2)  # Center in right half
    right_y = height - 150  # Near bottom
    
    # Draw text with outline
    for dx, dy in [(-2, -2), (-2, 2), (2, -2), (2, 2)]:
        draw.text((right_x + dx, right_y + dy), right_text, font=font_large, fill="black")
    # Main text
    draw.text((right_x, right_y), right_text, font=font_large, fill="white")
    
    # Brand text: "goodbald.com" (bottom center)
    brand_text = "goodbald.com"
    brand_bbox = draw.textbbox((0, 0), brand_text, font=font_brand)
    brand_width = brand_bbox[2] - brand_bbox[0]
    
    brand_x = (width // 2) - (brand_width // 2)  # Center horizontally
    brand_y = height - 50  # Bottom
    
    # Draw brand text with outline
    for dx, dy in [(-1, -1), (-1, 1), (1, -1), (1, 1)]:
        draw.text((brand_x + dx, brand_y + dy), brand_text, font=font_brand, fill="black")
    draw.text((brand_x, brand_y), brand_text, font=font_brand, fill="white")
    
    return img

def main():
    """Main execution function."""
    try:
        # Generate the base image
        base_img = generate_base_image()
        print(f"Base image generated: {base_img.size}")
        
        # Ensure it's the right size
        if base_img.size != (1080, 1080):
            print(f"Resizing from {base_img.size} to 1080x1080")
            base_img = base_img.resize((1080, 1080), Image.Resampling.LANCZOS)
        
        # Add text overlays
        final_img = add_text_overlays(base_img)
        print("Text overlays added")
        
        # Save the final image
        output_path = "/Users/igloo/.openclaw/workspace/scratch/creative-split-comparison.png"
        final_img.save(output_path, "PNG", quality=95)
        print(f"Final image saved to: {output_path}")
        
        # Verify success criteria
        print(f"Final image size: {final_img.size}")
        print(f"File exists: {os.path.exists(output_path)}")
        
        return True
        
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = main()
    if success:
        print("SUCCESS: Facebook ad image created successfully!")
    else:
        print("FAILED: Could not create image")