import gradio as gr
import requests
import os
from PIL import Image
import os
# Load environment variables
CHEVERETO_API_URL = os.environ.get("API_URL")  # Your Chevereto API endpoint
CHEVERETO_API_KEY = os.environ.get("API_KEY")  # Your Chevereto API Key
CHEVERETO_ALBUM_ID = os.environ.get("ALBUM_ID")  # Your Album ID

# Load the Gradio model
model = gr.Interface.load("models/stabilityai/stable-diffusion-3.5-large")

def upload_to_chevereto(img_path: str) -> str:
    """Uploads an image to Chevereto and returns the public image URL."""
    with open(img_path, "rb") as image_file:
        files = {"source": image_file}
        data = {
            "key": CHEVERETO_API_KEY,
            "format": "json",
            "album": CHEVERETO_ALBUM_ID  # Include the album ID
        }
        response = requests.post(CHEVERETO_API_URL, files=files, data=data)
    
    if response.status_code == 200:
        return response.json().get("image", {}).get("url")  # Return the image URL
    else:
        return f"Error uploading image: {response.text}"  # Debugging error message

def generate_image_and_upload(prompt: str) -> str:
    """Generates an image using the AI model and uploads it to Chevereto."""
    img = model(prompt)  # Generate image

    if isinstance(img, list):  
        img = img[0]  # If Gradio returns a list, get the first image

    if isinstance(img, Image.Image):
        img_path = "generated_image.png"
        img.save(img_path, format="PNG")  # Save locally before uploading

        return upload_to_chevereto(img_path)  # Upload and return URL
    else:
        return "Error: Model did not return a valid image"



# Define Gradio API returning the Chevereto public URL
iface = gr.Interface(fn=generate_image_and_upload, inputs="text", outputs="text")

iface.launch()