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

API_KEY = os.getenv("GOOGLE_API_KEY")

def image_generation(prompt):
    if not API_KEY:
        return "❌ API key not found. Please add GOOGLE_API_KEY in your Hugging Face secrets."

    # Use Gemini's free image generation model
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-preview-image-generation:generate?key={API_KEY}"
    headers = {"Content-Type": "application/json"}
    data = {"prompt": {"text": prompt}}

    try:
        res = requests.post(url, headers=headers, json=data)
    except Exception as e:
        return f"❌ Request failed: {str(e)}"

    if res.status_code != 200:
        return f"❌ Error {res.status_code}:\n{res.text}"

    try:
        result = res.json()
    except Exception:
        return f"❌ Invalid JSON response:\n{res.text}"

    # Extract image if present
    if "candidates" in result:
        for candidate in result["candidates"]:
            if "content" in candidate:
                for part in candidate["content"].get("parts", []):
                    if "inline_data" in part:
                        image_base64 = part["inline_data"].get("data")
                        if image_base64:
                            image = Image.open(BytesIO(base64.b64decode(image_base64)))
                            return image

    return f"⚠️ No image data found.\n\nFull response:\n{result}"

iface = gr.Interface(
    fn=image_generation,
    inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g. A futuristic city under sunset"),
    outputs=gr.Image(label="Generated Image"),
    title="Google Gemini Image Generator",
    description="Uses gemini-2.0-flash-preview-image-generation (free-tier). Educational purpose only."
)

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