import gradio as gr
from rembg import remove
from PIL import Image
import io

# Preserve original resolution & optimize for high-res images
def process_image(image, background_type, bg_color, bg_image):
    try:
        input_image = Image.open(image).convert("RGBA")
        # Use rembg with bytes instead of loading into RAM completely
        buffered = io.BytesIO()
        input_image.save(buffered, format="PNG")
        result = remove(buffered.getvalue())
        output = Image.open(io.BytesIO(result)).convert("RGBA")

        if background_type == "Solid Color":
            bg = Image.new("RGBA", output.size, bg_color)
            bg.paste(output, (0, 0), output)
            return bg

        elif background_type == "Custom Image" and bg_image:
            custom_bg = Image.open(bg_image).convert("RGBA").resize(output.size)
            custom_bg.paste(output, (0, 0), output)
            return custom_bg

        else:
            return output
    except Exception as e:
        return f"❌ Error: {str(e)}"

# For multiple images
def batch_process(images, background_type, bg_color, bg_image):
    return [process_image(img, background_type, bg_color, bg_image) for img in images]

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🎯 High-Resolution AI Background Remover with Custom Background")

    with gr.Row():
        images_input = gr.File(label="📤 Upload High-Res Image(s)", file_types=["image"], file_count="multiple")

    with gr.Row():
        background_type = gr.Radio(
            ["Transparent", "Solid Color", "Custom Image"],
            label="Choose Background Type",
            value="Transparent"
        )
        bg_color = gr.ColorPicker(label="Solid Color", value="#ffffff", visible=False)
        bg_image = gr.File(label="Custom Background Image", file_types=["image"], visible=False)

    output_gallery = gr.Gallery(label="🖼️ Output Images", columns=2, height=600)

    def toggle_visibility(bg_type):
        return (
            gr.update(visible=(bg_type == "Solid Color")),
            gr.update(visible=(bg_type == "Custom Image"))
        )

    background_type.change(toggle_visibility, inputs=background_type, outputs=[bg_color, bg_image])

    run_btn = gr.Button("🚀 Start Processing")
    run_btn.click(fn=batch_process, inputs=[images_input, background_type, bg_color, bg_image], outputs=output_gallery)

demo.launch()