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

API_KEY = os.environ.get("API_KEY")
API_URL = "https://api.4llm.com/v1/images/generations"

def generate_image(prompt, size="1024x1024"):
    start_time = time.time()  # Start timing
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    
    data = {
        "model": "4llm-image",
        "prompt": prompt,
        "n": 1,
        "size": size
    }
    
    try:
        response = requests.post(API_URL, headers=headers, json=data)
        response.raise_for_status()
        result = response.json()
        
        # Get the base64 image data
        image_data = result["data"][0]["url"].split(",")[1]
        image_bytes = base64.b64decode(image_data)
        
        # Convert to PIL Image
        image = Image.open(io.BytesIO(image_bytes))
        
        generation_time = time.time() - start_time  # Calculate generation time
        return image, result["data"][0]["revised_prompt"], f"{generation_time:.2f} seconds"
    except Exception as e:
        return None, f"Error: {str(e)}", ""

# Create the Gradio interface
with gr.Blocks(title="4LLM Image Generation", css="""
    footer { display: none !important; }
    .gr-button { background-color: #4CAF50; color: white; border: none; padding: 10px 20px; text-align: center; text-decoration: none; display: inline-block; font-size: 16px; margin: 4px 2px; cursor: pointer; border-radius: 5px; }
    .gr-button:hover { background-color: #45a049; }
    .gr-textbox { margin-bottom: 10px; }
    .gr-markdown { margin-bottom: 20px; }
""") as demo:
    gr.Markdown("# 4LLM Image Generation Demo")
    gr.Markdown("Generate images using the 4LLM API with no rate limits!")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image you want to generate...",
                lines=3
            )
            size = gr.Dropdown(
                choices=["1024x1024", "512x512", "768x768"],
                value="1024x1024",
                label="Image Size"
            )
            generate_btn = gr.Button("🔄 Generate Image")
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
            revised_prompt = gr.Textbox(
                label="Revised Prompt",
                lines=3,
                interactive=False
            )
            generation_time_display = gr.Textbox(
                label="Generation Time",
                lines=1,
                interactive=False
            )
            copy_btn = gr.Button("📋 Copy Image")
            download_btn = gr.Button("⬇️ Download Image")
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, size],
        outputs=[output_image, revised_prompt, generation_time_display]
    )

    # Add functionality for the Copy and Download buttons
    def copy_image(image):
        return image

    def download_image(image):
        return image

    copy_btn.click(fn=copy_image, inputs=output_image, outputs=None)
    download_btn.click(fn=download_image, inputs=output_image, outputs=None)

if __name__ == "__main__":
    demo.launch(share=False)