import gradio as gr
import numpy as np
from PIL import Image
import io
from huggingface_hub import InferenceClient

# Initialize client
client = InferenceClient()

def segment_and_edit(image, prompt):
    """
    Fast interior design editor using Hugging Face Inference API
    """
    try:
        # Convert PIL to bytes
        img_byte_arr = io.BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()
        
        # Use a free, public model for image-to-image
        edited_image = client.image_to_image(
            image=img_byte_arr,
            prompt=f"interior design photo, {prompt}, professional photography, high quality, detailed, realistic",
            negative_prompt="blurry, distorted, low quality, watermark, text, cartoon, sketch",
            model="black-forest-labs/FLUX.1-schnell",  # Fast, public model
        )
        
        return edited_image
        
    except Exception as e:
        # Fallback to text-to-image if image-to-image fails
        try:
            generated = client.text_to_image(
                prompt=f"interior room design, {prompt}, professional photography, 4k, realistic",
                model="black-forest-labs/FLUX.1-schnell"
            )
            return generated
        except Exception as e2:
            return None, f"Error: {str(e2)}"

def process_room(image, edit_prompt):
    """Main processing function"""
    if image is None:
        return None, "Please upload an image"
    
    if not edit_prompt.strip():
        return None, "Please provide editing instructions"
    
    # Process the image
    result = segment_and_edit(image, edit_prompt)
    
    if isinstance(result, tuple):
        return result
    
    return result, "✅ Edit complete!"

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏠 AI Interior Designer Assistant
    Upload a room photo and describe your desired changes. Uses AI to redesign your space.
    
    **Example prompts:**
    - "Change the sofa to mid-century modern, deep blue velvet"
    - "Replace with white minimalist furniture and add plants"
    - "Make it Scandinavian style with light wood"
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Room Photo")
            edit_prompt = gr.Textbox(
                label="Describe Your Changes",
                placeholder="E.g., change the sofa to modern gray fabric and add a floor lamp",
                lines=3
            )
            submit_btn = gr.Button("Transform Room 🎨", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Redesigned Room")
            status_text = gr.Textbox(label="Status", interactive=False)
    
    # Examples
    gr.Examples(
        examples=[
            ["Change to modern minimalist with white furniture", None],
            ["Replace with bohemian style, add colorful cushions and plants", None],
            ["Make it industrial loft style with exposed brick aesthetic", None],
        ],
        inputs=[edit_prompt, input_image],
    )
    
    submit_btn.click(
        fn=process_room,
        inputs=[input_image, edit_prompt],
        outputs=[output_image, status_text]
    )

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