import gradio as gr
from gradio_client import Client, handle_file
import re
import time
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Get Hugging Face token from environment variable
hf_token = os.getenv("HUGGING_FACE_HUB_TOKEN")

# Initialize client with auth
client = Client(
    "levihsu/OOTDiffusion",
    hf_token=hf_token
)


def generate_outfit(model_image, garment_image, n_samples=1, n_steps=20, image_scale=2, seed=-1):
    if model_image is None or garment_image is None:
        return None, "Please upload both model and garment images"
        
    max_retries = 3
    for attempt in range(max_retries):
        try:
            # Use the client to predict
            result = client.predict(
                vton_img=handle_file(model_image),
                garm_img=handle_file(garment_image),
                n_samples=n_samples,
                n_steps=n_steps,
                image_scale=image_scale,
                seed=seed,
                api_name="/process_hd"
            )
            
            # If result is a list, get the first item
            if isinstance(result, list):
                result = result[0]
            
            # If result is a dictionary, try to get the image path
            if isinstance(result, dict):
                if 'image' in result:
                    return result['image'], None
                else:
                    return None, "API returned unexpected format"
                
            return result, None
            
        except Exception as e:
            error_msg = str(e)
            if "exceeded your GPU quota" in error_msg:
                wait_time_match = re.search(r'retry in (\d+:\d+:\d+)', error_msg)
                wait_time = wait_time_match.group(1) if wait_time_match else "60:00"  # Default to 1 hour
                wait_seconds = sum(int(x) * 60 ** i for i, x in enumerate(reversed(wait_time.split(':'))))  # Convert wait time to seconds
                if attempt < max_retries - 1:
                    time.sleep(wait_seconds)  # Wait before retrying
                return None, f"GPU quota exceeded. Please wait {wait_time} before trying again."
            else:
                return None, f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("""
    ## Outfit Diffusion - Try On Virtual Outfits

    ⚠️ **Note**: This demo uses free GPU quota which is limited. To avoid errors:
    - Use lower values for Steps (10-15) and Scale (1-2)
    - Wait between attempts if you get a quota error
    - Sign up for a Hugging Face account for more quota

  
    """)

    with gr.Row():
        with gr.Column():
            model_image = gr.Image(
                label="Upload Model Image (person wearing clothes)", 
                type="filepath",
                height=300
                
            )
            model_examples = [
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/ba5ba7978e7302e8ab5eb733cc7221394c4e6faf/model_5.png",
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/40dade4a04a827c0fdf63c6c70b42ef26480f391/01861_00.jpg",
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/3c4639c5fab3cdcd3239609dca5afee7b0677286/model_6.png",
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/0089171df270f4532eec3d80a8f36cc8218c6840/01008_00.jpg"
            ]
            gr.Examples(examples=model_examples, inputs=model_image)

            garment_image = gr.Image(
                label="Upload Garment Image (clothing item)", 
                type="filepath",
                height=300
            )
            garment_examples = [
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/180d4e2a1139071a8685a5edee7ab24bcf1639f5/03244_00.jpg",
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/584dda2c5ee1d8271a6cd06225c07db89c79ca03/04825_00.jpg",
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/a51938ec99f13e548d365a9ca6d794b6fe7462af/049949_1.jpg",
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/2d64241101189251ce415df84dc9205cda9a36ca/03032_00.jpg",
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/44aee6b576cae51eeb979311306375b56b7e0d8b/02305_00.jpg",
                "https://levihsu-ootdiffusion.hf.space/file=/tmp/gradio/578dfa869dedb649e91eccbe566fc76435bb6bbe/049920_1.jpg"
            ]
            gr.Examples(examples=garment_examples, inputs=garment_image)

        
        with gr.Column():
            output_image = gr.Image(label="Generated Output")
            error_text = gr.Markdown()  # Add error display
    
    with gr.Row():
        with gr.Column():
            n_samples = gr.Slider(
                label="Number of Samples", 
                minimum=1, 
                maximum=5, 
                step=1, 
                value=1
            )
            n_steps = gr.Slider(
                label="Steps (lower = faster, try 10-15)", 
                minimum=1, 
                maximum=50, 
                step=1, 
                value=10  # Reduced default
            )
            image_scale = gr.Slider(
                label="Scale (lower = faster, try 1-2)", 
                minimum=1, 
                maximum=5, 
                step=1, 
                value=1  # Reduced default
            )
            seed = gr.Number(
                label="Random Seed (-1 for random)", 
                value=-1
            )
    
    generate_button = gr.Button("Generate Outfit")

    # Set up the action for the button
    generate_button.click(
        fn=generate_outfit,
        inputs=[model_image, garment_image, n_samples, n_steps, image_scale, seed],
        outputs=[output_image, error_text]
    )

# Launch the app
demo.launch()