#!/usr/bin/env python3
"""
TEST - Imagen 3 API (Image Generation)
"""
import os
import requests
import base64
from PIL import Image
import io

API_KEY = os.getenv('GEMINI_API_KEY')
print(f"✅ API Key: {API_KEY[:15]}...")

# Use Imagen 3 endpoint
API_URL = "https://generativelanguage.googleapis.com/v1beta/models/imagen-3.0-generate-001:predict"

prompt = "Dark hand-drawn 2D stylized indie game art. Character KAI - male, age 17, green and pink dreadlocks, tactical survivor outfit. Walking toward camera, front view. Bold black outlines, centered, transparent background."

print(f"\n🎨 Generating with Imagen 3...")
print(f"Prompt: {prompt[:80]}...\n")

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

headers = {"Content-Type": "application/json"}
url = f"{API_URL}?key={API_KEY}"

print(f"📡 Calling Imagen API...")
try:
    response = requests.post(url, headers=headers, json=payload, timeout=60)
    print(f"Status: {response.status_code}")
    
    if response.status_code == 200:
        data = response.json()
        print(f"✅ Response received!")
        print(f"Response keys: {list(data.keys())}")
        
        if 'predictions' in data and len(data['predictions']) > 0:
            pred = data['predictions'][0]
            print(f"Prediction keys: {list(pred.keys())}")
            
            if 'bytesBase64Encoded' in pred:
                print(f"✅ Found image data!")
                image_b64 = pred['bytesBase64Encoded']
                image_bytes = base64.b64decode(image_b64)
                image = Image.open(io.BytesIO(image_bytes))
                
                output_path = "assets/slike/TEST_kai_imagen.png"
                image.save(output_path, "PNG")
                print(f"✅✅ SAVED: {output_path}")
                print(f"Image size: {image.size}")
        else:
            print(f"Response data: {data}")
    else:
        print(f"❌ Error {response.status_code}: {response.text}")
        
except Exception as e:
    print(f"❌ Exception: {e}")

print("\n✅ Test complete!")