import gradio as gr
import numpy as np
from PIL import Image


def get_map(p):
    d = np.array(p.convert("RGB")).astype(float)
    return np.max(d, axis=2) / 255.0


def rm_wm(img, m, x, y, strength):
    h, w = m.shape
    roi = img[y : y + h, x : x + w].astype(float)
    m = np.clip(m * strength, 0, 0.99)
    for i in range(3):
        roi[:, :, i] = (roi[:, :, i] - m * 255) / (1 - m)
    img[y : y + h, x : x + w] = np.clip(roi, 0, 255).astype(np.uint8)
    return img


def clean(src, strength):
    try:
        if src is None:
            return None, None

        mask48 = Image.open("assets/bg_48.png")
        mask96 = Image.open("assets/bg_96.png")

        arr = np.array(src)
        w, h = src.size
        is_big = w > 1024 and h > 1024
        sz, mg = (96, 64) if is_big else (48, 32)
        mask = get_map(mask96 if is_big else mask48)
        x, y = w - mg - sz, h - mg - sz
        return src, Image.fromarray(rm_wm(arr, mask, x, y, strength))
    except FileNotFoundError:
        gr.Warning("Mask files not found in assets folder")
        return src, src
    except Exception as e:
        gr.Warning(f"Error: {str(e)}")
        return src, src


css = """
* {
    font-family: 'Cascadia Code', 'Consolas', 'Monaco', monospace !important;
}

.gradio-container {
    max-width: 1000px !important;
    margin: auto !important;
}

#title {
    font-size: 24px !important;
    font-weight: 600 !important;
    color: #1a1a1a !important;
    margin-bottom: 8px !important;
}

#description {
    color: #666 !important;
    font-size: 14px !important;
    margin-bottom: 24px !important;
}

.gr-button-primary {
    background: #2563eb !important;
    border: none !important;
    color: white !important;
    font-weight: 500 !important;
}

.gr-button-primary:hover {
    background: #1d4ed8 !important;
}

input[type="range"] {
    accent-color: #2563eb !important;
}

.image-container {
    border: 1px solid #e5e7eb !important;
    border-radius: 8px !important;
    overflow: hidden !important;
}
"""

demo = gr.Interface(
    fn=clean,
    inputs=[
        gr.Image(type="pil", label="Source Image"),
        gr.Slider(minimum=0, maximum=2, value=1, step=0.1, label="Removal Strength"),
    ],
    outputs=gr.ImageSlider(label="Original <-> Cleaned"),
    title="Gemini Watermark Remover",
    description="Credits: [Formula original](https://github.com/journey-ad/gemini-watermark-remover) // Source code: [Github](https://github.com/oopshnik/gemini)",
    css=css,
)

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