from __future__ import annotations

import argparse
import shutil
from pathlib import Path

import numpy as np
from PIL import Image, ImageFilter


def inpaint_pixels(data: np.ndarray, mask: np.ndarray) -> np.ndarray:
    result = data.astype(np.float32).copy()
    remaining = mask.copy()

    for _ in range(data.shape[0] + data.shape[1]):
        known = ~remaining
        count = np.zeros(remaining.shape, dtype=np.float32)
        total = np.zeros(result.shape, dtype=np.float32)

        neighbor = known[:-1, :]
        count[1:, :] += neighbor
        total[1:, :] += result[:-1, :] * neighbor[..., None]

        neighbor = known[1:, :]
        count[:-1, :] += neighbor
        total[:-1, :] += result[1:, :] * neighbor[..., None]

        neighbor = known[:, :-1]
        count[:, 1:] += neighbor
        total[:, 1:] += result[:, :-1] * neighbor[..., None]

        neighbor = known[:, 1:]
        count[:, :-1] += neighbor
        total[:, :-1] += result[:, 1:] * neighbor[..., None]

        boundary = remaining & (count > 0)
        if not boundary.any():
            break

        result[boundary] = total[boundary] / count[boundary, None]
        remaining[boundary] = False

        if not remaining.any():
            break

    if remaining.any():
        known_pixels = result[~remaining]
        fill = np.median(known_pixels, axis=0) if len(known_pixels) else np.array([255, 255, 255])
        result[remaining] = fill

    return np.clip(result, 0, 255).astype(np.uint8)


def build_watermark_mask(image: Image.Image) -> Image.Image:
    rgb = image.convert("RGB")
    data = np.asarray(rgb).astype(np.int16)
    height, width = data.shape[:2]

    yy, xx = np.mgrid[:height, :width]
    light_right_width = min(width, max(140, min(190, int(width * 0.27))))
    dark_right_width = min(width, max(180, min(360, int(width * 0.42))))
    light_bottom_height = min(height, max(28, min(72, int(height * 0.1))))
    dark_bottom_height = min(height, max(28, min(54, int(height * 0.075))))
    light_right_band = xx >= width - light_right_width
    dark_right_band = xx >= width - dark_right_width
    light_target_corner = light_right_band & (yy >= height - light_bottom_height)
    dark_target_corner = dark_right_band & (yy >= height - dark_bottom_height)

    channel_max = data.max(axis=2)
    channel_min = data.min(axis=2)
    saturation = channel_max - channel_min
    brightness = data.mean(axis=2)
    blurred = np.asarray(rgb.filter(ImageFilter.GaussianBlur(3))).astype(np.int16)
    contrast = np.abs(data - blurred).mean(axis=2)

    corner_brightness = brightness[light_target_corner]
    local_background = float(np.median(corner_brightness)) if corner_brightness.size else 255.0

    light_background_mark = (
        (saturation <= 36)
        & (brightness >= 165)
        & (brightness <= 254)
        & (channel_min >= 160)
    )

    dark_background_mark = (
        (saturation <= 48)
        & (brightness >= max(58, local_background + 14))
        & (brightness <= 240)
        & (channel_min >= 42)
    )

    if local_background < 155:
        contrast_mark = (
            (contrast >= 5)
            & (saturation <= 105)
            & (brightness >= max(70, local_background + 16))
            & (brightness <= 252)
        )
        mask = dark_target_corner & (dark_background_mark | contrast_mark)
    else:
        contrast_mark = (
            (contrast >= 4)
            & (saturation <= 72)
            & (brightness <= local_background - 3)
            & (brightness >= 145)
        )
        mask = light_target_corner & (light_background_mark | contrast_mark)

    mask_img = Image.fromarray((mask.astype(np.uint8)) * 255, mode="L")
    mask_img = mask_img.filter(ImageFilter.MaxFilter(3))
    mask_img = mask_img.filter(ImageFilter.GaussianBlur(0.6))
    return mask_img


def remove_watermark(source: Path, destination: Path, quality: int) -> bool:
    with Image.open(source) as opened:
        image = opened.convert("RGBA")

    mask = build_watermark_mask(image)
    if not mask.getbbox():
        destination.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(source, destination)
        return False

    base = image.convert("RGB")
    data = np.asarray(base).copy()
    soft_mask = np.asarray(mask)
    mask_array = soft_mask > 12

    if mask_array.any():
        ys, xs = np.where(mask_array)
        pad = 20
        y0 = max(0, int(ys.min()) - pad)
        y1 = min(data.shape[0], int(ys.max()) + pad + 1)
        x0 = max(0, int(xs.min()) - pad)
        x1 = min(data.shape[1], int(xs.max()) + pad + 1)
        roi = data[y0:y1, x0:x1]
        roi_mask = mask_array[y0:y1, x0:x1]
        inpainted = inpaint_pixels(roi, roi_mask)
        data[y0:y1, x0:x1][roi_mask] = inpainted[roi_mask]

    filled = Image.fromarray(data, mode="RGB")
    cleaned = Image.composite(filled.convert("RGBA"), image, mask)

    destination.parent.mkdir(parents=True, exist_ok=True)
    cleaned.save(destination, format="WEBP", quality=quality, method=6)
    return True


def main() -> None:
    parser = argparse.ArgumentParser(description="Remove article platform watermarks from lower-right corners.")
    parser.add_argument("--source", default="public/articles")
    parser.add_argument("--out", default=".tmp/watermark-cleaned")
    parser.add_argument("--quality", type=int, default=96)
    args = parser.parse_args()

    source_root = Path(args.source)
    out_root = Path(args.out)
    files = sorted(source_root.rglob("*.awebp"))
    changed = 0

    for source in files:
        destination = out_root / source.relative_to(source_root)
        if remove_watermark(source, destination, args.quality):
            changed += 1

    print(f"processed={len(files)} changed={changed} out={out_root}")


if __name__ == "__main__":
    main()