import gradio as gr
import sys
import logging
from inference import main, load_index, load_metadata
from PIL import Image

def run_pipeline(prompt, image):
    """
    Gradio interface function to run the inference pipeline
    """
    try:
        logging.info("Loading required data...")
        index = load_index()
        metadata_df = load_metadata()
        
        logging.info("Starting inference pipeline...")
        results = main(prompt, image, index, metadata_df)
        
        # Return the generated image and similar products
        image_path = results['generated_image_path']
        similar_products = results['similar_products']

        image_output = Image.open(image_path)
        
        # Format product URLs as a numbered list with similarity scores
        product_urls = []
        for i, product in enumerate(similar_products, 1):
            similarity = 1 / (1 + product['distance'])
            product_urls.append(f"{i}. Similarity: {similarity:.2f}\nProduct: {product['product_url']}\n")
        
        formatted_urls = "\n".join(product_urls)
        return image_output, formatted_urls
        
    except Exception as e:
        logging.error(f"Error in pipeline: {str(e)}")
        return None, None

# Create Gradio interface
iface = gr.Interface(
    fn=run_pipeline,
    inputs=[
        gr.Textbox(label="Enter your prompt", placeholder="e.g., modern living room with minimalist furniture"),
        gr.Image(label="Upload control image", type="filepath")
    ],
    outputs=[
        gr.Image(label="Generated Image"),
        gr.Textbox(label="Similar IKEA Products", lines=15)
    ],
    title="Interior Design Image Generator",
    description="Upload an image and provide a prompt to generate interior design variations and find similar IKEA products.",
    theme="default",
    flagging_mode="never"
)

if __name__ == "__main__":
    iface.launch(share=True)