import gradio as gr
import requests
import os
import time

# Environment variables for API details
API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")  # Fetching the API token from environment variable

# Function to query Hugging Face API
def query_huggingface_api(api_url, prompt):
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    data = {"inputs": prompt}
    response = requests.post(api_url, headers=headers, json=data)
    
    if response.status_code == 200:
        return response.content, None  # Return the image and no error
    else:
        return None, f"Error {response.status_code}: {response.text}"  # Return None and the error message

# Gradio function for generating the image
def generate_image(api_url, prompt):
    # Attempt to query the API with retry logic for loading models
    for attempt in range(5):  # Try up to 5 times
        result, error = query_huggingface_api(f"https://api-inference.huggingface.co/models/{api_url}", prompt)
        
        if result:
            return result, None
        elif "Model is currently loading" in error:
            estimated_time = float(error.split("estimated_time\":")[1].split("}")[0])  # Extract estimated time from error message
            time.sleep(estimated_time + 5)  # Wait for the model to load, with an additional buffer time
        else:
            return None, error  # Return the error if it's not a loading issue

    return None, "Model is still loading after multiple attempts. Please try again later."  # Final error if all attempts fail

# Create Gradio Blocks Interface
with gr.Blocks(theme="nevreal/blues") as demo:
    gr.Markdown(
        """
        # Text to Image Generator
        Enter a text prompt, and the custom model will generate an image.
        """
    )

    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Enter your prompt", 
                placeholder="Type something here...",
                value=""
            )
            model_input = gr.Textbox(
                label="Model URL", 
                placeholder="Enter the model URL...",
                value=""
            )
            generate_btn = gr.Button("Generate Image")
        
        with gr.Column():
            image_output = gr.Image(label="Generated Image")
            error_output = gr.Textbox(label="Error", interactive=False)

    # Define the action for the button
    generate_btn.click(
        fn=generate_image, 
        inputs=[model_input, text_input],  # Pass both model URL and prompt
        outputs=[image_output, error_output]
    )

# Launch the Gradio Blocks WebUI
demo.launch(share=True, debug=True)