# app.py
# A user-friendly version with a custom download button, inspired by the watermark remover app.

import gradio as gr
from rembg import remove
import traceback
from PIL import Image
import tempfile # Used to create temporary files for downloading

# --- Core Function ---
# This function processes the image. It remains the same.
def process_image(input_image):
    """
    Removes the background from an image.
    Handles errors gracefully.
    """
    if input_image is None:
        raise gr.Error("Please upload an image or select one of the examples.")
    
    try:
        # Use the default rembg model for stability
        return remove(input_image)
    except Exception as e:
        print(traceback.format_exc())
        raise gr.Error(f"Failed to process image: {e}")

# --- Download Function ---
# This function correctly prepares the image for download.
def prepare_for_download(image):
    """
    Saves the processed PIL image to a temporary file and returns the path.
    This is the correct way to trigger a download in Gradio.
    """
    if image is None:
        # Prevent an error if the download button is clicked before an image is processed.
        return None
    try:
        # Create a temporary file with a '.png' extension.
        # 'delete=False' is important so the file is not deleted before Gradio serves it.
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
            image.save(tmp.name)
            # Return the path to the temporary file
            return tmp.name
    except Exception as e:
        print(f"Error creating download file: {e}")
        return None

# --- Gradio App with Modern UI ---
# We inject CSS to hide the default utility buttons on the output image.
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="violet"),
    css="#output_image .meta-buttons { display: none !important; }"
) as demo:
    
    gr.Markdown(
        """
        Upload an image, use your webcam, or try one of the examples below to instantly remove the background.
        """
    )
    
    # Main two-column layout
    with gr.Row(variant="panel"):
        
        # --- Input Column ---
        with gr.Column(scale=1, min_width=300):
            image_input = gr.Image(
                type="pil",
                label="Upload or Drag an Image",
                sources=["upload", "webcam"], 
                height=400
            )
            
           
            with gr.Row():
                clear_btn = gr.ClearButton(
                    value="Clear",
                    components=[image_input], 
                    variant="secondary"
                )
                submit_btn = gr.Button("Remove Background", variant="primary")

        # --- Output Column ---
        with gr.Column(scale=1, min_width=300):
            # We give the output image an ID so we can target it with CSS
            image_output = gr.Image(
                elem_id="output_image",
                type="pil",
                label="Result",
                height=400,
                interactive=False # This disables the upload overlay on the result
            )
            
            # ✅ NEW: Use a gr.File component for the download, which is more reliable.
            # This will appear as a download link/button when populated.
            download_output = gr.File(label="Download Result", visible=False)

    # --- Event Listeners ---
    def on_submit(input_img):
        # This function runs when the user clicks submit.
        # It processes the image and then prepares it for download.
        processed_img = process_image(input_img)
        download_path = prepare_for_download(processed_img)
        # It returns the processed image to the preview box,
        # and the downloadable file to the (now visible) download component.
        return processed_img, gr.update(value=download_path, visible=True)

    submit_btn.click(
        fn=on_submit,
        inputs=image_input,
        outputs=[image_output, download_output]
    )

    # When the user clicks "Clear", also hide the download button.
    clear_btn.click(lambda: gr.update(visible=False), outputs=[download_output])

# --- Launch the App ---
if __name__ == "__main__":
    demo.launch()