import gradio as gr
import numpy as np
import cv2
from PIL import Image
import tempfile
import os

def dilate_mask(mask, strength=5):
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (strength, strength))
    return cv2.dilate(mask, kernel, iterations=1)

def patch_based_inpaint(img_bgr, mask, dilate_strength=5):
    """
    Patch Clone: untuk watermark kecil, ambil patch dari sekitar dan tempel.
    Jauh lebih tajam daripada cv2.inpaint yang blur.
    """
    mask = dilate_mask(mask, dilate_strength)
    h, w = img_bgr.shape[:2]
    
    # Bounding box watermark
    x, y, mw, mh = cv2.boundingRect(mask)
    if mw == 0 or mh == 0:
        return img_bgr
    
    # Perluas bbox sedikit untuk blending
    pad = max(mw, mh) // 2
    x1 = max(0, x - pad)
    y1 = max(0, y - pad)
    x2 = min(w, x + mw + pad)
    y2 = min(h, y + mh + pad)
    
    # Cari patch sumber: ambil dari area di sekitar bbox yang TIDAK kena mask
    # Prioritas: kiri -> atas -> kanan -> bawah (arah yang paling jauh dari tepi gambar)
    patch_h = y2 - y1
    patch_w = x2 - x1
    
    candidates = []
    
    # Kiri
    if x1 - patch_w >= 0:
        patch = img_bgr[y1:y2, x1-patch_w:x1].copy()
        if patch.shape[0] == patch_h and patch.shape[1] == patch_w:
            candidates.append(patch)
    
    # Atas
    if y1 - patch_h >= 0:
        patch = img_bgr[y1-patch_h:y1, x1:x2].copy()
        if patch.shape[0] == patch_h and patch.shape[1] == patch_w:
            candidates.append(patch)
    
    # Kanan
    if x2 + patch_w <= w:
        patch = img_bgr[y1:y2, x2:x2+patch_w].copy()
        if patch.shape[0] == patch_h and patch.shape[1] == patch_w:
            candidates.append(patch)
    
    # Bawah
    if y2 + patch_h <= h:
        patch = img_bgr[y2:y2+patch_h, x1:x2].copy()
        if patch.shape[0] == patch_h and patch.shape[1] == patch_w:
            candidates.append(patch)
    
    if not candidates:
        # Fallback: cv2.inpaint NS
        return cv2.inpaint(img_bgr, mask, 3, cv2.INPAINT_NS)
    
    # Pilih patch yang paling mirip dengan tepi area watermark (L2 distance minimal)
    # Ambil ring 5px di sekitar bbox sebagai referensi
    best_patch = candidates[0]
    if len(candidates) > 1:
        ring_mask = np.zeros((h, w), dtype=np.uint8)
        cv2.rectangle(ring_mask, (x1, y1), (x2, y2), 255, 3)
        ring_mask = cv2.bitwise_and(ring_mask, cv2.bitwise_not(mask))
        
        if cv2.countNonZero(ring_mask) > 0:
            ref_mean = cv2.mean(img_bgr, ring_mask)[:3]
            best_dist = float('inf')
            for cand in candidates:
                # Resize candidate ke ukuran bbox untuk compare
                cand_small = cv2.resize(cand, (x2-x1, y2-y1))
                # Ambil border candidate
                cm = np.zeros((y2-y1, x2-x1), dtype=np.uint8)
                cv2.rectangle(cm, (0,0), (x2-x1-1, y2-y1-1), 255, 3)
                if cv2.countNonZero(cm) > 0:
                    dist = sum(abs(a-b) for a,b in zip(cv2.mean(cand_small, cm)[:3], ref_mean))
                    if dist < best_dist:
                        best_dist = dist
                        best_patch = cand
    
    # Resize patch ke ukuran bbox
    patch_resized = cv2.resize(best_patch, (x2-x1, y2-y1), interpolation=cv2.INTER_CUBIC)
    
    # Buat mask untuk blending (feathered edge)
    roi_mask = mask[y1:y2, x1:x2].copy()
    if roi_mask.shape[0] != patch_resized.shape[0] or roi_mask.shape[1] != patch_resized.shape[1]:
        roi_mask = cv2.resize(roi_mask, (patch_resized.shape[1], patch_resized.shape[0]), interpolation=cv2.INTER_NEAREST)
    
    # Feather mask untuk blending halus
    roi_mask_f = roi_mask.astype(np.float32) / 255.0
    roi_mask_f = cv2.GaussianBlur(roi_mask_f, (0,0), sigmaX=max(3, dilate_strength//2))
    
    # Blend patch ke gambar asli
    result = img_bgr.copy()
    roi = result[y1:y2, x1:x2].astype(np.float32)
    patch_f = patch_resized.astype(np.float32)
    
    for c in range(3):
        roi[:,:,c] = roi[:,:,c] * (1 - roi_mask_f) + patch_f[:,:,c] * roi_mask_f
    
    result[y1:y2, x1:x2] = roi.astype(np.uint8)
    
    # Post: slight sharpening untuk area hasil
    kernel_sharp = np.array([[0,-1,0],[-1,5,-1],[0,-1,0]], dtype=np.float32)
    sharpened = cv2.filter2D(result, -1, kernel_sharp)
    blend_sharp = roi_mask_f * 0.3
    for c in range(3):
        result[y1:y2, x1:x2, c] = result[y1:y2, x1:x2, c] * (1-blend_sharp) + sharpened[y1:y2, x1:x2, c] * blend_sharp
    
    return result

def ns_sharpen_inpaint(img_bgr, mask, dilate_strength=2):
    """
    Navier-Stokes inpaint + unsharp mask + CLAHE.
    Lebih baik untuk tekstur metal/kompleks daripada Telea.
    """
    mask = dilate_mask(mask, dilate_strength)
    
    # NS lebih tajam untuk edge/tekstur daripada Telea
    result = cv2.inpaint(img_bgr, mask, 3, cv2.INPAINT_NS)
    
    # Unsharp mask untuk mengurangi blur
    gaussian = cv2.GaussianBlur(result, (0,0), sigmaX=2.0)
    result = cv2.addWeighted(result, 1.5, gaussian, -0.5, 0)
    
    # CLAHE untuk kontras lokal (membantu tekstur metal)
    lab = cv2.cvtColor(result, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    l = clahe.apply(l)
    result = cv2.cvtColor(cv2.merge([l,a,b]), cv2.COLOR_LAB2BGR)
    
    return result

def remove_watermark(editor_data, method_name, dilate_strength, detail_strength):
    if editor_data is None:
        return None, None, None
    
    try:
        if not isinstance(editor_data, dict):
            return None, None, None
        
        bg = editor_data.get("background")
        layers = editor_data.get("layers", [])
        
        if bg is None:
            return None, None, None
        
        # Parse background
        if isinstance(bg, np.ndarray):
            img = Image.fromarray(bg).convert("RGB")
        elif hasattr(bg, "convert"):
            img = bg.convert("RGB")
        else:
            img = Image.fromarray(np.array(bg)).convert("RGB")
        
        if len(layers) == 0:
            return img, None, None
        
        # Parse mask dari layer brush
        mask_layer = layers[-1]
        mask = np.array(mask_layer) if not isinstance(mask_layer, np.ndarray) else mask_layer
        
        img_np = np.array(img)
        img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
        
        if len(mask.shape) == 3:
            if mask.shape[2] == 4:
                mask_gray = cv2.cvtColor(mask, cv2.COLOR_RGBA2GRAY)
            else:
                mask_gray = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
        else:
            mask_gray = mask
        
        _, mask_bin = cv2.threshold(mask_gray, 5, 255, cv2.THRESH_BINARY)
        
        if cv2.countNonZero(mask_bin) == 0:
            return img, None, None
        
        # Pilih metode
        if method_name == "Patch Clone (Tajam)":
            result_bgr = patch_based_inpaint(img_bgr, mask_bin, dilate_strength)
        elif method_name == "NS + Sharpen (Smooth)":
            result_bgr = ns_sharpen_inpaint(img_bgr, mask_bin, dilate_strength)
        else:
            # Telea original
            mask_d = dilate_mask(mask_bin, dilate_strength)
            result_bgr = cv2.inpaint(img_bgr, mask_d, 3, cv2.INPAINT_TELEA)
        
        result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)
        result_pil = Image.fromarray(result_rgb)
        
        # Simpan ke file temporary untuk download
        png_path = os.path.join(tempfile.gettempdir(), "result.png")
        jpg_path = os.path.join(tempfile.gettempdir(), "result.jpg")
        
        result_pil.save(png_path, format="PNG")
        result_pil.save(jpg_path, format="JPEG", quality=95, optimize=True)
        
        return result_pil, png_path, jpg_path
        
    except Exception as e:
        # Fallback aman
        try:
            if isinstance(editor_data, dict) and editor_data.get("background"):
                bg = editor_data["background"]
                if isinstance(bg, np.ndarray):
                    fallback = Image.fromarray(bg).convert("RGB")
                else:
                    fallback = bg.convert("RGB") if hasattr(bg, "convert") else None
                return fallback, None, None
        except Exception:
            return None, None, None

# ==========================================
# GRADIO 6.x UI
# ==========================================

with gr.Blocks(title="Watermark Remover CPU - Patch Clone") as demo:
    gr.Markdown("""
    # 🧽 Watermark Remover (Patch Clone Edition)
    ### CPU-Friendly untuk Hugging Face Space Gratis
    
    **Cara pakai:**
    1. Upload gambar di canvas (ikon 📎 pojok kanan bawah)
    2. Pilih ikon **kuas (brush)** dan coret area watermark
    3. Pilih metode:
       - **Patch Clone (Tajam)** → ambil tekstur dari sekitar, tidak blur
       - **NS + Sharpen (Smooth)** → Navier-Stokes + sharpening, blur minimal
    4. Atur **Dilate Mask** (2-5 untuk watermark kecil, 10+ untuk besar)
    5. Klik **✨ Hapus Watermark**
    6. Download **PNG** atau **JPG**
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            editor = gr.ImageEditor(
                label="📸 Upload & Brush Area Watermark",
                height=420,
            )
            
            with gr.Row():
                method = gr.Radio(
                    ["Patch Clone (Tajam)", "NS + Sharpen (Smooth)", "Telea (Blur)"],
                    value="Patch Clone (Tajam)",
                    label="Metode"
                )
            
            dilate = gr.Slider(
                minimum=1, maximum=25, value=3, step=1,
                label="🔧 Dilate Mask (2-5 untuk watermark kecil)"
            )
            
            detail = gr.Slider(
                minimum=0, maximum=100, value=80, step=5,
                label="🎨 Detail Strength (unused, reserved)",
                visible=False  # hidden, reserved for future
            )
            
            btn_remove = gr.Button("✨ Hapus Watermark", variant="primary", size="lg")
            btn_clear = gr.Button("🗑️ Reset")
            
        with gr.Column(scale=1):
            output_img = gr.Image(
                label="✅ Hasil Preview", 
                height=350,
            )
            
            gr.Markdown("### 📥 Download Hasil")
            with gr.Row():
                output_png = gr.File(label="Download PNG", interactive=False)
                output_jpg = gr.File(label="Download JPG", interactive=False)
            
            gr.Markdown("""
            **Kenapa tidak blur lagi?**
            - 🧩 **Patch Clone** mengambil tekstur asli dari sekitar watermark dan menempelkannya dengan blending halus
            - ⚡ **NS + Sharpen** menggunakan Navier-Stokes (lebih tajam dari Telea) + unsharp mask + CLAHE
            - 🎯 Untuk watermark kecil di pojok/pojok: **Patch Clone** adalah solusi terbaik untuk CPU
            - 💾 Output tersedia dalam **PNG (lossless)** dan **JPG (quality 95)**
            """)
    
    btn_remove.click(
        fn=remove_watermark,
        inputs=[editor, method, dilate, detail],
        outputs=[output_img, output_png, output_jpg]
    )
    
    btn_clear.click(
        lambda: (None, None, None),
        outputs=[output_img, output_png, output_jpg]
    )

if __name__ == "__main__":
    demo.launch()