from __future__ import annotations

import cv2
import numpy as np

from .models import DocumentSpec, MaskStroke, Point, TextRegion
from .shape_recognition import shape_source_mask


def build_mask(rgba: np.ndarray, document: DocumentSpec) -> np.ndarray:
    height, width = rgba.shape[:2]
    combined = np.zeros((height, width), dtype=np.uint8)
    for region in document.regions:
        if region.remove_text:
            combined = cv2.bitwise_or(combined, region_mask(rgba, region))
    for shape in document.shapes:
        if shape.remove_source:
            shape_mask = shape_source_mask(shape, width, height)
            padding = max(0, int(round(shape.mask_padding)))
            if padding:
                kernel = cv2.getStructuringElement(
                    cv2.MORPH_ELLIPSE,
                    (padding * 2 + 1, padding * 2 + 1),
                )
                shape_mask = cv2.dilate(shape_mask, kernel)
            combined = cv2.bitwise_or(combined, shape_mask)
    for stroke in document.mask_strokes:
        apply_stroke(combined, stroke)
    return combined


def region_mask(rgba: np.ndarray, region: TextRegion) -> np.ndarray:
    height, width = rgba.shape[:2]
    polygon = np.array([[point.x, point.y] for point in region.source_quad], dtype=np.int32)
    bounds = cv2.boundingRect(polygon)
    x, y, w, h = bounds
    x0, y0 = max(0, x), max(0, y)
    x1, y1 = min(width, x + w), min(height, y + h)
    output = np.zeros((height, width), dtype=np.uint8)
    if x1 <= x0 or y1 <= y0:
        return output

    crop = rgba[y0:y1, x0:x1, :3]
    lab = cv2.cvtColor(crop, cv2.COLOR_RGB2LAB).astype(np.float32)
    local_polygon = polygon - np.array([x0, y0])
    polygon_mask = np.zeros(crop.shape[:2], dtype=np.uint8)
    cv2.fillPoly(polygon_mask, [local_polygon], 255)

    perimeter = cv2.morphologyEx(polygon_mask, cv2.MORPH_GRADIENT, np.ones((3, 3), np.uint8))
    border_pixels = lab[perimeter > 0]
    background = np.median(border_pixels, axis=0) if border_pixels.size else np.median(lab.reshape(-1, 3), axis=0)
    delta = np.linalg.norm(lab - background, axis=2)
    inside = delta[polygon_mask > 0]
    if inside.size:
        otsu_threshold, _ = cv2.threshold(
            np.clip(inside, 0, 255).astype(np.uint8),
            0,
            255,
            cv2.THRESH_BINARY + cv2.THRESH_OTSU,
        )
        threshold = max(8.0, float(otsu_threshold) * 0.72)
    else:
        threshold = 8.0
    foreground = ((delta >= threshold) & (polygon_mask > 0)).astype(np.uint8) * 255

    count, labels, stats, _ = cv2.connectedComponentsWithStats(foreground, 8)
    filtered = np.zeros_like(foreground)
    max_area = max(9, int(w * h * 0.22))
    min_area = 1
    for index in range(1, count):
        area = stats[index, cv2.CC_STAT_AREA]
        component_height = stats[index, cv2.CC_STAT_HEIGHT]
        if min_area <= area <= max_area and component_height >= 1:
            filtered[labels == index] = 255
    if cv2.countNonZero(filtered) < max(3, w * h * 0.005):
        filtered = foreground

    padding = max(1, int(round(region.mask_padding)))
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (padding * 2 + 1, padding * 2 + 1))
    filtered = cv2.dilate(filtered, kernel)
    output[y0:y1, x0:x1] = filtered
    return output


def apply_stroke(mask: np.ndarray, stroke: MaskStroke) -> None:
    value = 255 if stroke.operation == "add" else 0
    radius = max(1, int(round(stroke.radius)))
    points = [(int(round(point.x)), int(round(point.y))) for point in stroke.points]
    if len(points) == 1:
        cv2.circle(mask, points[0], radius, value, -1, lineType=cv2.LINE_AA)
        return
    for start, end in zip(points, points[1:]):
        cv2.line(mask, start, end, value, radius * 2, lineType=cv2.LINE_AA)
        cv2.circle(mask, end, radius, value, -1, lineType=cv2.LINE_AA)