import os
import gradio as gr
from dotenv import load_dotenv

load_dotenv()

style_transfer_function = gr.load(
    "prgarg007/ai-style-transfer-dreamsofa",
    src="spaces",
    token=os.getenv("HF_TOKEN")
)

def process_with_private_space(swatch_image, product_image, room_image, room_style):
    """
    Call the private space as a function and return the results
    """
    if style_transfer_function is None:
        return None, None, "Error: Could not connect to the AI processing service. Please try again later."
    
    if not swatch_image or not product_image or not room_image:
        return None, None, "Please upload all three images (swatch, product, and room)."
    
    try:
        # Call the private space as a function using the API name we defined
        styled_product, final_result, status_message = style_transfer_function(
            swatch_image, 
            product_image, 
            room_image, 
            room_style,
            api_name="process_style_transfer"
        )
        
        return styled_product, final_result, status_message
        
    except Exception as e:
        return None, None, f"Error processing images: {str(e)}"

def create_public_interface():
    """Create the public interface that uses the private space as a function"""
    
    with gr.Blocks(
        title="AI Style Transfer & Room Integration - Public Access",
        theme=gr.themes.Soft()
    ) as app:
        
        gr.Markdown("""
        # 🚀 AI Powered Style Transfer & Room Integration
        ## Public Access Portal
        
        Transform your furniture with AI-powered style transfer and room integration!
        
        **How it works:**
        1. Upload a **swatch image** (fabric/pattern you want to apply)
        2. Upload a **product image** (furniture piece to style)
        3. Upload a **room image** (where you want to place the furniture)
        4. Select your preferred **room style**
        5. Let AI work its magic! ✨
        """)
        
        # Check if the private space is loaded
        if style_transfer_function is None:
            gr.Markdown("""
            ⚠️ **Service Currently Unavailable**
            
            The AI processing service is currently unavailable. This might be due to:
            - Network connectivity issues
            - Authentication problems
            - Service maintenance
            
            Please try again later or contact support.
            """)
            return app
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📤 Upload Your Images")
                
                swatch_input = gr.Image(
                    label="🎨 Swatch/Pattern Image",
                    type="pil",
                    height=200
                )
                
                product_input = gr.Image(
                    label="📦 Furniture/Product Image", 
                    type="pil",
                    height=200
                )
                
                room_input = gr.Image(
                    label="🏠 Room/Background Image",
                    type="pil", 
                    height=200,
                )
                
                room_style = gr.Dropdown(
                    choices=["modern", "vintage", "industrial", "scandinavian", "bohemian", "rustic"],
                    value="modern",
                    label="🎭 Room Style Theme",
                )
                
                process_btn = gr.Button(
                    "🤖 Transform with AI", 
                    variant="primary", 
                    size="lg",
                    scale=1
                )
                
                gr.Markdown("""
                ### 💡 Tips for Best Results:
                - Use high-quality, clear images
                - Ensure good lighting in your photos
                - Choose images with similar perspectives
                - Swatch images work best when the pattern/texture is clearly visible
                """)
            
            with gr.Column(scale=1):
                gr.Markdown("### 🎯 AI Generated Results")
                
                styled_product_output = gr.Image(
                    label="🎨 Step 1: Styled Product",
                    height=300
                )
                
                final_output = gr.Image(
                    label="🏆 Step 2: Room Integration", 
                    height=300
                )
                
                status_output = gr.Textbox(
                    label="🔄 Processing Status",
                    value="Ready to transform your furniture with AI...",
                    interactive=False
                )
                
                gr.Markdown("""
                ### 🚀 Processing Steps:
                1. **AI Analysis**: Examines your swatch and product images
                2. **Style Transfer**: Applies the swatch style to your furniture  
                3. **Room Integration**: Places the styled furniture in your room
                4. **Final Rendering**: Creates photorealistic results
                
                *Processing typically takes 30-60 seconds*
                """)
        
        # Connect the button to our function that calls the private space
        process_btn.click(
            fn=process_with_private_space,
            inputs=[swatch_input, product_input, room_input, room_style],
            outputs=[styled_product_output, final_output, status_output],
            show_progress=True
        )
        
        # Add some example information
        gr.Markdown("""
        ---
        ### 🎨 Example Use Cases:
        - **Interior Design**: Preview how different fabrics look on your furniture
        - **E-commerce**: Show products in various styles and room settings  
        - **Home Renovation**: Experiment with different design aesthetics
        - **Furniture Customization**: Visualize custom upholstery options
        
        ### ⚡ Powered by:
        - OpenAI DALL-E for image generation
        - Advanced AI for style analysis and transfer
        - Gradio for seamless user experience
        
        ### 🔒 Privacy & Security:
        - Your images are processed securely
        - No images are permanently stored
        - Processing happens in real-time
        """)
    
    return app

if __name__ == "__main__":
    demo = create_public_interface()
    demo.launch()