import gradio as gr
from together import Together
from PIL import Image
import requests
from io import BytesIO
import os
from datetime import datetime

# Retrieve API key from Hugging Face Secrets
API_KEY = os.getenv("MY_API_KEY")

client = Together(api_key=API_KEY)

# Ensure directory for saving user-generated images
IMAGE_FOLDER = "./user_images"
if not os.path.exists(IMAGE_FOLDER):
    os.makedirs(IMAGE_FOLDER, exist_ok=True)
    print(f"Created folder: {os.path.abspath(IMAGE_FOLDER)}")

# Function to generate an image and save it
def generate_and_save_image(prompt):
    try:
        if not prompt.strip():
            return None, None  # Avoid empty prompts
        
        print(f"Generating image for: {prompt}")

        # Generate image with Together AI API
        response = client.images.generate(
            prompt=prompt,
            model="black-forest-labs/FLUX.1-schnell",  # Model name
            steps=4
        )
        
        # Fetch the image from the returned URL
        image_url = response.data[0].url
        print(f"Image URL: {image_url}")
        response_image = requests.get(image_url)
        image = Image.open(BytesIO(response_image.content))

        # Save the image locally
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        image_path = os.path.join(IMAGE_FOLDER, f"{timestamp}.png")
        image.save(image_path)
        print(f"Image successfully saved at: {os.path.abspath(image_path)}")

        return image, image_path  # Return the generated image and its path
    except Exception as e:
        print("Error:", e)
        return None, None

# Function to get all saved images for the gallery
def get_gallery_images():
    try:
        images = []
        for filename in sorted(os.listdir(IMAGE_FOLDER)):
            if filename.endswith(".png"):
                images.append(os.path.join(IMAGE_FOLDER, filename))
        print(f"Gallery images: {images}")
        return images
    except Exception as e:
        print("Error fetching gallery images:", e)
        return []

# Gradio Interface Setup
with gr.Blocks() as app:
    gr.Markdown("## Real-Time AI Image Generator 🚀")
    gr.Markdown("Start typing a prompt below to generate images in real-time with Together AI's FLUX model. Your creations will be displayed below in the gallery.")

    with gr.Row():
        prompt_box = gr.Textbox(label="Enter your prompt", placeholder="A horse on the moon")
        image_output = gr.Image(label="Generated Image")

    # Add a gallery section
    with gr.Row():
        gallery = gr.Gallery(label="Generated with Love by You", height="auto")
    
    # Automatically trigger image generation and save image
    def process_and_update_gallery(prompt):
        image, image_path = generate_and_save_image(prompt)
        if image:
            # Return the generated image and updated gallery
            return image, get_gallery_images()
        return None, get_gallery_images()
    
    prompt_box.change(
        fn=process_and_update_gallery,
        inputs=prompt_box,
        outputs=[image_output, gallery]
    )

# Launch the Gradio app
app.launch()