import gradio as gr
from rembg import remove
from PIL import Image
import numpy as np
import time
import requests
from io import BytesIO

def washing_machine(image):
    # Convert to PIL Image if it's a numpy array
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # Simulate washing process with progress updates
    for i in range(1, 6):
        time.sleep(0.5)  # Simulate washing cycles
        yield f"Washing cycle {i}/5...", image
    
    # Actual background removal
    result = remove(image)
    
    # Add some "spin cycle" effects
    for angle in [15, -15, 0]:
        rotated = result.rotate(angle)
        yield "Spin cycle complete!", rotated
    
    # Final clean image
    yield "Your image is clean and background-free!", result

def load_image_from_url(url):
    try:
        response = requests.get(url)
        img = Image.open(BytesIO(response.content))
        return img
    except Exception as e:
        raise gr.Error(f"Failed to load image from URL: {str(e)}")

# Custom CSS for washing machine theme
css = """
#wash-btn {background-color: #00a8ff; color: white; border: none; padding: 10px 20px; text-align: center; text-decoration: none; display: inline-block; font-size: 16px; margin: 4px 2px; cursor: pointer; border-radius: 12px;}
#wash-btn:hover {background-color: #0095e6;}
.title {font-family: Arial; color: #333; text-align: center;}
.description {font-family: Arial; color: #666; text-align: center;}
#output-image {border: 5px dashed #00a8ff; border-radius: 10px;}
#input-image {border: 5px dashed #ccc; border-radius: 10px;}
.tab {padding: 20px;}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("<h1 class='title'>🚀 Background Washing Machine 🧼</h1>")
    gr.Markdown("<p class='description'>Upload your image and we'll wash the background right off!</p>")
    
    with gr.Tabs():
        with gr.TabItem("Upload Image"):
            with gr.Row():
                with gr.Column():
                    file_input = gr.Image(label="Dirty Image (with background)", type="pil")
                    file_btn = gr.Button("Start Washing!", elem_id="wash-btn")
                
                with gr.Column():
                    file_status = gr.Textbox(label="Washing Status")
                    file_output = gr.Image(label="Clean Image (no background)")
        
        with gr.TabItem("From URL"):
            with gr.Row():
                with gr.Column():
                    url_input = gr.Textbox(label="Image URL", placeholder="https://example.com/image.jpg")
                    url_fetch = gr.Button("Fetch Image", elem_id="wash-btn")
                    url_btn = gr.Button("Start Washing!", elem_id="wash-btn", visible=False)
                    url_preview = gr.Image(label="Preview", interactive=False)
                
                with gr.Column():
                    url_status = gr.Textbox(label="Washing Status")
                    url_output = gr.Image(label="Clean Image (no background)")
    
    gr.Markdown("### How it works:")
    gr.Markdown("1. Upload your image 🖼️ or paste a URL\n2. Click the washing button 🧼\n3. Wait while we wash away the background 🫧\n4. Get your clean, background-free image! ✅")

    # File upload tab functionality
    file_btn.click(
        washing_machine,
        inputs=file_input,
        outputs=[file_status, file_output]
    )
    
    # URL tab functionality
    def fetch_image(url):
        if not url:
            raise gr.Error("Please enter a valid URL")
        img = load_image_from_url(url)
        return img, gr.Button(visible=True)
    
    url_fetch.click(
        fetch_image,
        inputs=url_input,
        outputs=[url_preview, url_btn]
    )
    
    url_btn.click(
        washing_machine,
        inputs=url_preview,
        outputs=[url_status, url_output]
    )

demo.launch()