#!/usr/bin/env python3
import os
import sys
import requests
import base64
import time
from pathlib import Path

# Imagen API configuration
API_KEY = "AIzaSyB-nR0nkeftKrd42NrNIDcFCj3yFP8JLtw"
MODEL = "imagen-4.0-generate-preview-06-06"
API_URL = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:predict"

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

# Plugin configurations
plugins = [
    {
        'filename': 'daily-appointment-summary.png',
        'prompt': '''A modern, minimalist app icon for "Daily Appointment Summary Email" plugin.
Square format with rounded corners. Indigo blue gradient background (#4f46e5).
White icon showing a combination of an envelope and a calendar.
Flat design, clean lines, professional SaaS application style.
The icon should be simple and recognizable at small sizes.'''
    },
    {
        'filename': 'no-show-tracker.png',
        'prompt': '''A modern, minimalist app icon for "No-Show Customer Tracker" plugin.
Square format with rounded corners. Red gradient background (#dc2626).
White icon showing a person silhouette with an X or cancel symbol overlay.
Flat design, clean lines, professional SaaS application style.
The icon should convey missed appointments clearly at small sizes.'''
    },
    {
        'filename': 'birthday-greetings.png',
        'prompt': '''A modern, minimalist app icon for "Birthday Greeting Campaign" plugin.
Square format with rounded corners. Pink gradient background (#ec4899).
White icon showing a birthday cake or gift box with a bow.
Flat design, clean lines, cheerful yet professional style.
The icon should be celebratory and friendly at small sizes.'''
    },
    {
        'filename': 'monthly-revenue-report.png',
        'prompt': '''A modern, minimalist app icon for "Monthly Revenue Report" plugin.
Square format with rounded corners. Green gradient background (#10b981).
White icon showing an upward trending chart or bar graph.
Flat design, clean lines, professional business analytics style.
The icon should convey growth and success at small sizes.'''
    },
    {
        'filename': 'appointment-reminder-24hr.png',
        'prompt': '''A modern, minimalist app icon for "Appointment Reminder" plugin.
Square format with rounded corners. Amber/orange gradient background (#f59e0b).
White icon showing a notification bell with a small red badge or alert dot.
Flat design, clean lines, attention-grabbing yet professional style.
The icon should convey alerts and reminders clearly at small sizes.'''
    },
    {
        'filename': 'inactive-customer-reengagement.png',
        'prompt': '''A modern, minimalist app icon for "Inactive Customer Re-engagement" plugin.
Square format with rounded corners. Purple gradient background (#8b5cf6).
White icon showing a heart with a circular refresh arrow around it.
Flat design, clean lines, warm and welcoming professional style.
The icon should convey customer care and return at small sizes.'''
    }
]

def generate_image(prompt, filename):
    """Generate an image using Imagen API"""
    print(f"\nGenerating {filename}...")

    headers = {
        "Content-Type": "application/json"
    }

    payload = {
        "instances": [{
            "prompt": prompt
        }],
        "parameters": {
            "sampleCount": 1
        }
    }

    try:
        response = requests.post(
            f"{API_URL}?key={API_KEY}",
            headers=headers,
            json=payload,
            timeout=120
        )

        if response.status_code == 200:
            result = response.json()

            # Check if there's an image in the predictions
            if 'predictions' in result and len(result['predictions']) > 0:
                prediction = result['predictions'][0]
                if 'bytesBase64Encoded' in prediction:
                    # Extract and save the image
                    image_data = base64.b64decode(prediction['bytesBase64Encoded'])
                    output_path = os.path.join(OUTPUT_DIR, filename)
                    with open(output_path, 'wb') as f:
                        f.write(image_data)
                    print(f"✓ Saved: {output_path} ({len(image_data)} bytes)")
                    return True
                elif 'image' in prediction and 'bytesBase64Encoded' in prediction['image']:
                    image_data = base64.b64decode(prediction['image']['bytesBase64Encoded'])
                    output_path = os.path.join(OUTPUT_DIR, filename)
                    with open(output_path, 'wb') as f:
                        f.write(image_data)
                    print(f"✓ Saved: {output_path} ({len(image_data)} bytes)")
                    return True

            print(f"✗ No image generated. Response: {str(result)[:300]}...")
            return False
        else:
            print(f"✗ API Error ({response.status_code}): {response.text[:500]}")
            return False

    except Exception as e:
        print(f"✗ Error: {str(e)}")
        return False

def main():
    # Create output directory
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("Starting logo generation with Imagen 4.0 API...")
    print("=" * 60)

    success_count = 0
    for i, plugin in enumerate(plugins):
        if i > 0:
            # Add delay between requests to avoid rate limiting
            print("Waiting 2 seconds before next request...")
            time.sleep(2)

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

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

if __name__ == "__main__":
    main()