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

# ========== تحميل الموديل ==========
_model = None
_device = None

def get_model():
    global _model, _device
    if _model is not None:
        return _model, _device

    _device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[LaMa] Loading on {_device}...")

    # جديد ✅
    model_path = hf_hub_download(
        repo_id="fashn-ai/LaMa",
        filename="big-lama.pt",
    )
    

    _model = torch.jit.load(model_path, map_location=_device)
    _model.eval()
    print("[LaMa] Ready.")
    return _model, _device


def _pad(tensor, multiple=8):
    h, w = tensor.shape[-2:]
    ph = (multiple - h % multiple) % multiple
    pw = (multiple - w % multiple) % multiple
    if ph or pw:
        tensor = torch.nn.functional.pad(tensor, (0, pw, 0, ph), mode="reflect")
    return tensor, h, w


def lama_inpaint(image_np, mask_np):
    model, device = get_model()

    img  = image_np.astype(np.float32) / 255.0
    mask = (mask_np > 127).astype(np.float32)

    img_t  = torch.from_numpy(img).permute(2,0,1).unsqueeze(0).to(device)
    mask_t = torch.from_numpy(mask).unsqueeze(0).unsqueeze(0).to(device)

    masked = img_t * (1 - mask_t)

    masked_p, h, w = _pad(masked)
    mask_p,   _,  _ = _pad(mask_t)

    with torch.no_grad():
        try:
            out = model(masked_p, mask_p)
        except Exception:
            inp = torch.cat([masked_p, mask_p], dim=1)
            out = model(inp)

    out = out[:, :, :h, :w]
    result = img_t[:, :, :h, :w] * (1 - mask_t) + out * mask_t
    result = result.squeeze(0).permute(1,2,0).cpu().numpy()
    return (result * 255).clip(0, 255).astype(np.uint8)


# ========== الواجهة ==========
def process(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_np = np.array(bg.convert("RGB"))

    if not layers or layers[0] is None:
        return bg.convert("RGB")

    alpha = np.array(layers[0].convert("RGBA"))[:, :, 3]

    if feather > 0:
        alpha = np.array(
            Image.fromarray(alpha).filter(ImageFilter.GaussianBlur(radius=feather))
        )

    result = lama_inpaint(image_np, alpha)
    return Image.fromarray(result)


with gr.Blocks(title="SFX Cleaner - LaMa") as demo:
    gr.Markdown("""
    # 🧹 SFX Cleaner — big-lama
    **الطريقة:** ارفع الصورة ← ارسم بالفرشاة البيضاء على النص/SFX ← اضغط **تبييض**
    > أول تشغيل يحمل الموديل تلقائياً (~200MB)
    """)

    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="نعومة الحواف")
        btn = gr.Button("🧹 تبييض", variant="primary", size="lg")

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

demo.launch()