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

# Your endpoint URL and token (you might want to use environment variables for these in production)
API_URL = "https://ca80xvp8jmhqanbz.us-east-1.aws.endpoints.huggingface.cloud"

def process_image(input_image, prompt):
    # Convert Gradio's image input to base64
    if input_image is None:
        return None, "Please upload an image"
    
    # Convert PIL Image directly to base64 string
    buffer = io.BytesIO()
    input_image.save(buffer, format="WEBP")
    base64_image = f"data:image/webp;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
    
    # Prepare headers and data
    headers = {
        "Authorization": f"Bearer {os.getenv('HF_API_KEY')}",
        "Content-Type": "application/json",
        "Accept": "application/json"
    }
    
    data = {
        "inputs": {
            "prompt": prompt,
            "image": base64_image
        }
    }
    
    # Make the API request with retry mechanism
    max_retries = 5
    retry_delay = 60  # seconds
    
    for retry in range(max_retries):
        try:
            response = requests.post(API_URL, headers=headers, json=data, timeout=300)
            
            if response.status_code == 200:
                # Convert the response base64 directly to PIL Image
                result = response.json()
                image_bytes = base64.b64decode(result["final_image"])
                output_image = Image.open(io.BytesIO(image_bytes))
                
                # Return both the image and the parameters used
                params_text = "\nParameters used:\n"
                for key, value in result["parameters"].items():
                    params_text += f"{key}: {value}\n"
                
                return output_image, params_text
                
            elif response.status_code == 503:
                if retry < max_retries - 1:
                    time.sleep(retry_delay)
                    continue
                return None, f"Service unavailable after {max_retries} retries"
            else:
                return None, f"Error: {response.status_code}\n{response.text}"
                
        except requests.exceptions.RequestException as e:
            if retry < max_retries - 1:
                time.sleep(retry_delay)
                continue
            return None, f"Request error: {str(e)}"
    
    return None, "Maximum retries reached"

# Create Gradio interface
with gr.Blocks(title="Room Design Diffusion", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # Room Design Diffusion
    Upload a room image and provide a prompt describing how you'd like to redesign it.
    The AI will generate a new version of your room based on your description.
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Upload Room Image", type="pil")
            prompt = gr.Textbox(
                label="Describe how you want to redesign the room",
                placeholder="Example: A modern living room with a comfortable gray sectional sofa, glass coffee table, minimalist TV stand, and geometric area rug",
                lines=3
            )
            submit_btn = gr.Button("Generate Design", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Generated Design")
            output_text = gr.Textbox(label="Processing Details", lines=5)
    
    submit_btn.click(
        fn=process_image,
        inputs=[input_image, prompt],
        outputs=[output_image, output_text]
    )
    
    gr.Markdown("""
    ### Tips for best results:
    - Use clear, detailed prompts describing the desired room style
    - Make sure the input image is well-lit and clearly shows the room
    - Be patient as generation can take a few minutes
    """)

# Launch the app with sharing enabled
demo.launch(share=True)