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

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def generate_image(prompt):
    try:
        RUNPOD_API_KEY = os.getenv("RUNPOD_API_KEY")
        RUNPOD_ENDPOINT = os.getenv("RUNPOD_ENDPOINT")
        
        if not all([RUNPOD_API_KEY, RUNPOD_ENDPOINT]):
            return None, "API credentials not configured"
        
        payload = {
            "input": {
                "prompt": prompt,
                "width": 768,
                "height": 768,
                "num_inference_steps": 20,
                "return_base64": True
            }
        }
        
        headers = {
            "Authorization": f"Bearer {RUNPOD_API_KEY}",
            "Content-Type": "application/json"
        }
        
        logger.info(f"Sending request to RunPod: {payload}")
        response = requests.post(RUNPOD_ENDPOINT, json=payload, headers=headers, timeout=30)
        
        if response.status_code != 200:
            return None, f"API Error {response.status_code}: {response.text}"
        
        data = response.json()
        
        if "image" in data:
            image_data = base64.b64decode(data["image"].split(",")[1])
            return Image.open(BytesIO(image_data)), "Success!"
        elif "output" in data and data["output"]:
            image_url = data["output"][0]
            img_data = requests.get(image_url, timeout=30).content
            return Image.open(BytesIO(img_data)), "Success!"
        else:
            return None, "Unexpected response format"
            
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        return None, f"Error: {str(e)}"

with gr.Blocks() as app:
    gr.Markdown("# 🎨 RunPod Image Generator")
    
    with gr.Row():
        prompt = gr.Textbox(label="Enter prompt", value="A comic book superhero")
        generate_btn = gr.Button("Generate")
    
    output = gr.Image(label="Generated Image")
    status = gr.Textbox(label="Status")
    
    generate_btn.click(
        generate_image,
        inputs=[prompt],
        outputs=[output, status]
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)