import gradio as gr
import numpy as np
from PIL import Image, ImageFilter
import torch
import cv2

device = "cuda" if torch.cuda.is_available() else "cpu"

# تحميل lama_large_512px — نفس الموديل اللي في BallonsTranslator
from iopaint.model.lama import LaMa
from iopaint.schema import InpaintRequest, HDStrategy

model = LaMa(device=device)

def inpaint(editor_data, feather):
    if editor_data is None:
        return None

    bg = editor_data.get("background")
    layers = editor_data.get("layers", [])

    if bg is None:
        return None

    image = bg.convert("RGB")
    image_np = np.array(image)

    if not layers or layers[0] is None:
        return image

    # استخراج الماسك
    layer = layers[0].convert("RGBA")
    mask_np = np.array(layer)[:, :, 3]

    # feather اختياري
    if feather > 0:
        mask_pil = Image.fromarray(mask_np)
        mask_pil = mask_pil.filter(ImageFilter.GaussianBlur(radius=feather))
        mask_np = np.array(mask_pil)

    mask_np = (mask_np > 127).astype(np.uint8) * 255

    config = InpaintRequest(hd_strategy=HDStrategy.ORIGINAL)
    result = model(image_np, mask_np, config)

    return Image.fromarray(result)


with gr.Blocks(title="SFX Cleaner - LaMa Large") as demo:
    gr.Markdown("""
    # 🧹 SFX Cleaner — lama_large_512px
    **الطريقة:** ارفع الصورة ← ارسم بالفرشاة على النص/SFX ← اضغط **تبييض**
    """)

    with gr.Row():
        with gr.Column(scale=2):
            editor = gr.ImageEditor(
                label="📌 ارسم على النص",
                brush=gr.Brush(
                    colors=["#ffffff"],
                    color_mode="fixed",
                    default_size=20
                ),
                eraser=gr.Eraser(default_size=20),
                type="pil",
                height=650,
            )
        with gr.Column(scale=1):
            output = gr.Image(
                label="✅ النتيجة",
                type="pil",
                height=650,
            )

    with gr.Row():
        feather = gr.Slider(0, 8, value=2, step=1, label="نعومة الحواف (Feather)")
        btn = gr.Button("🧹 تبييض", variant="primary", size="lg")

    btn.click(fn=inpaint, inputs=[editor, feather], outputs=output)

demo.launch()