# coding=utf-8
"""Turn a gr.ImageEditor value into a background image, a stroke mask and a pixel box."""
from __future__ import annotations

import cv2
import numpy as np
from PIL import Image

# Diff threshold for the fallback path; the frontend re-encodes images as webp,
# so compression noise has to be filtered out.
DIFF_THRESHOLD = 32
MIN_COMPONENT_AREA = 24


def keep_large_components(mask, min_area=MIN_COMPONENT_AREA):
    """Drop small blobs caused by compression noise."""
    m = np.asarray(mask).astype(np.uint8)
    n, labels, stats, _ = cv2.connectedComponentsWithStats(m, connectivity=8)
    out = np.zeros(m.shape, dtype=bool)
    for i in range(1, n):
        if stats[i, cv2.CC_STAT_AREA] >= min_area:
            out |= labels == i
    return out


def _layer_alpha_mask(layer, size):
    if layer is None:
        return None
    if layer.size != size:
        layer = layer.resize(size, Image.BILINEAR)
    return np.asarray(layer.convert("RGBA"))[..., 3] > 10


def editor_mask(editor_value):
    """Return (background RGB, stroke mask).

    Prefer the brush layer alpha; fall back to a background/composite diff when the
    frontend does not send the layer.
    """
    if not editor_value:
        return None, None
    if not isinstance(editor_value, dict):
        bg = editor_value.convert("RGB")
        return bg, np.zeros((bg.size[1], bg.size[0]), dtype=bool)
    bg = editor_value.get("background")
    if bg is None:
        return None, None
    bg = bg.convert("RGB")
    h, w = bg.size[1], bg.size[0]

    layer_mask = np.zeros((h, w), dtype=bool)
    for layer in (editor_value.get("layers") or []):
        m = _layer_alpha_mask(layer, bg.size)
        if m is not None:
            layer_mask |= m
    if layer_mask.any():
        return bg, keep_large_components(layer_mask)

    comp = editor_value.get("composite")
    if comp is None:
        return bg, np.zeros((h, w), dtype=bool)
    comp = comp.convert("RGB")
    if comp.size != bg.size:
        comp = comp.resize(bg.size, Image.BILINEAR)
    diff = np.abs(np.asarray(bg, np.int16) - np.asarray(comp, np.int16)).max(axis=2)
    return bg, keep_large_components(diff > DIFF_THRESHOLD)


def mask_to_box(mask, pad_px=0):
    """Bounding box of the mask as (x0, y0, x1, y1), right/bottom exclusive."""
    ys, xs = np.nonzero(np.asarray(mask))
    if len(xs) == 0:
        return None
    pad = int(max(0, pad_px))
    return (int(xs.min()) - pad, int(ys.min()) - pad,
            int(xs.max()) + 1 + pad, int(ys.max()) + 1 + pad)


def clamp_box(box, img_w, img_h):
    """Clamp a box to the image bounds; None if it is degenerate."""
    if box is None:
        return None
    x0, y0, x1, y1 = [float(v) for v in box]
    if x1 < x0:
        x0, x1 = x1, x0
    if y1 < y0:
        y0, y1 = y1, y0
    x0, x1 = max(0.0, min(x0, img_w)), max(0.0, min(x1, img_w))
    y0, y1 = max(0.0, min(y0, img_h)), max(0.0, min(y1, img_h))
    if x1 - x0 < 2 or y1 - y0 < 2:
        return None
    return (int(round(x0)), int(round(y0)), int(round(x1)), int(round(y1)))


def manual_box(values, img_w, img_h):
    """Parse manual x0/y0/x1/y1 inputs; None when invalid."""
    try:
        x0, y0, x1, y1 = [float(v) for v in values]
    except (TypeError, ValueError):
        return None
    if x1 <= x0 or y1 <= y0:
        return None
    return clamp_box((x0, y0, x1, y1), img_w, img_h)