#!/usr/bin/env python3
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import os
import time

# Setup Client
client = genai.Client(api_key="AIzaSyB-nR0nkeftKrd42NrNIDcFCj3yFP8JLtw")

OUTPUT_DIR = "/home/poduck/Desktop/smoothschedule2/frontend/public/plugin-logos"

# Plugin configurations with detailed prompts
plugins = [
    {
        'filename': 'daily-appointment-summary.png',
        'prompt': '''Create a modern, minimalist app icon in square format with rounded corners.
Design: Indigo blue gradient background (#4f46e5 to lighter blue).
Icon: Simple white envelope overlaid with a small calendar icon in the corner.
Style: Flat design, clean geometric shapes, professional SaaS aesthetic.
The icon should be crisp and recognizable at 48x48 pixels.'''
    },
    {
        'filename': 'no-show-tracker.png',
        'prompt': '''Create a modern, minimalist app icon in square format with rounded corners.
Design: Red gradient background (#dc2626 to darker red).
Icon: Simple white person silhouette with a bold white X symbol overlaid.
Style: Flat design, clean geometric shapes, professional warning aesthetic.
The icon should clearly convey "missed appointment" at small sizes.'''
    },
    {
        'filename': 'birthday-greetings.png',
        'prompt': '''Create a modern, minimalist app icon in square format with rounded corners.
Design: Pink gradient background (#ec4899 to lighter pink).
Icon: Simple white birthday cake with 3 candles or a white gift box with a bow.
Style: Flat design, clean geometric shapes, cheerful yet professional aesthetic.
The icon should feel celebratory and friendly at small sizes.'''
    },
    {
        'filename': 'monthly-revenue-report.png',
        'prompt': '''Create a modern, minimalist app icon in square format with rounded corners.
Design: Green gradient background (#10b981 to darker green).
Icon: Simple white upward trending line chart or ascending bar graph.
Style: Flat design, clean geometric shapes, professional business analytics aesthetic.
The icon should convey growth and success at small sizes.'''
    },
    {
        'filename': 'appointment-reminder-24hr.png',
        'prompt': '''Create a modern, minimalist app icon in square format with rounded corners.
Design: Amber/orange gradient background (#f59e0b to darker orange).
Icon: Simple white notification bell with a small red circular alert badge.
Style: Flat design, clean geometric shapes, attention-grabbing yet professional aesthetic.
The icon should clearly indicate alerts and reminders at small sizes.'''
    },
    {
        'filename': 'inactive-customer-reengagement.png',
        'prompt': '''Create a modern, minimalist app icon in square format with rounded corners.
Design: Purple gradient background (#8b5cf6 to darker purple).
Icon: Simple white heart symbol with a circular white refresh/return arrow around it.
Style: Flat design, clean geometric shapes, warm and welcoming professional aesthetic.
The icon should convey customer care and returning customers at small sizes.'''
    }
]

def generate_logo(prompt, filename):
    """Generate a logo using Gemini 2.5 Flash Image"""
    print(f"\nGenerating {filename}...")

    try:
        response = client.models.generate_content(
            model='gemini-2.5-flash-image',
            contents=prompt,
            config=types.GenerateContentConfig(
                response_modalities=["IMAGE"]
            )
        )

        # Save the Image
        if response.candidates[0].content.parts:
            for part in response.candidates[0].content.parts:
                if part.inline_data:
                    image_data = part.inline_data.data
                    image = Image.open(BytesIO(image_data))

                    # Save to output directory
                    output_path = os.path.join(OUTPUT_DIR, filename)
                    image.save(output_path)
                    print(f"✓ Saved: {output_path}")
                    return True
                else:
                    print(f"✗ Model returned text instead of image: {part.text[:100]}")
                    return False
        else:
            print("✗ No content parts in response")
            return False

    except Exception as e:
        error_msg = str(e)
        if "RESOURCE_EXHAUSTED" in error_msg or "quota" in error_msg.lower():
            print(f"✗ Quota exceeded. Please try again tomorrow when quota resets.")
            print(f"   Error: {error_msg[:200]}...")
        else:
            print(f"✗ Error: {type(e).__name__}: {error_msg[:200]}")
        return False

def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("=" * 70)
    print("Generating Plugin Logos with Gemini 2.5 Flash Image")
    print("=" * 70)

    success_count = 0
    for i, plugin in enumerate(plugins):
        if i > 0:
            # Wait 3 seconds between requests to avoid rate limiting
            print("\nWaiting 3 seconds before next request...")
            time.sleep(3)

        if generate_logo(plugin['prompt'], plugin['filename']):
            success_count += 1

    print("\n" + "=" * 70)
    print(f"Generation complete: {success_count}/{len(plugins)} successful")
    print("=" * 70)

    if success_count == 0:
        print("\nNote: If you hit quota limits, try again tomorrow when the quota resets!")

if __name__ == "__main__":
    main()