"""Locate the visible Gemini sparkle and build a mask for shared inpainting."""

# OpenCV and NumPy expose incomplete types at this array-processing boundary.
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalCall=false, reportOptionalSubscript=false, reportOptionalOperand=false, reportAttributeAccessIssue=false, reportPrivateImportUsage=false, reportPrivateUsage=false, reportInvalidTypeForm=false, reportConstantRedefinition=false, reportUnnecessaryComparison=false
from __future__ import annotations

import functools
import logging
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, Any

import cv2
import numpy as np

from remove_ai_watermarks import image_io

if TYPE_CHECKING:
    from collections.abc import Iterator

    from numpy.typing import NDArray

logger = logging.getLogger(__name__)


class WatermarkSize(Enum):
    """Provider size tier selected from the source dimensions."""

    SMALL = "small"
    LARGE = "large"


@dataclass
class DetectionResult:
    """Detection decision and its component scores."""

    detected: bool = False
    confidence: float = 0.0
    region: tuple[int, int, int, int] = (0, 0, 0, 0)
    size: WatermarkSize = WatermarkSize.SMALL
    spatial_score: float = 0.0
    gradient_score: float = 0.0
    variance_score: float = 0.0


@dataclass(frozen=True, slots=True)
class WatermarkPosition:
    """Expected provider margins and logo size."""

    margin_right: int
    margin_bottom: int
    logo_size: int

    def get_position(self, image_width: int, image_height: int) -> tuple[int, int]:
        return image_width - self.margin_right - self.logo_size, image_height - self.margin_bottom - self.logo_size


@dataclass(frozen=True, slots=True)
class _Candidate:
    scale: int
    x: int
    y: int
    spatial: float
    gradient: float = 0.0
    variance: float = 0.0

    @property
    def fused(self) -> float:
        if self.spatial < 0.25:
            return max(0.0, self.spatial * 0.5)
        return self.spatial * 0.50 + self.gradient * 0.30 + self.variance * 0.20


def get_watermark_size(width: int, height: int) -> WatermarkSize:
    """Return the provider's large tier only when both axes exceed 1024."""
    return WatermarkSize.LARGE if width > 1024 and height > 1024 else WatermarkSize.SMALL


def get_watermark_config(width: int, height: int) -> WatermarkPosition:
    """Return the observed standard placement for the selected size tier."""
    if get_watermark_size(width, height) is WatermarkSize.LARGE:
        return WatermarkPosition(64, 64, 96)
    return WatermarkPosition(32, 32, 48)


def _calculate_alpha_map(background_capture: NDArray[Any]) -> NDArray[Any]:
    """Convert a black-background sparkle capture to a normalized opacity map."""
    if background_capture.ndim == 2:
        intensity = background_capture
    elif background_capture.shape[2] >= 3:
        intensity = background_capture[:, :, :3].max(axis=2)
    else:
        intensity = background_capture[:, :, 0]
    return intensity.astype(np.float32) / 255.0


def _load_capture(filename: str, expected_side: int) -> NDArray[Any]:
    capture = image_io.imread(Path(__file__).parent / "assets" / filename, cv2.IMREAD_COLOR)
    if capture is None:
        raise RuntimeError(f"Failed to decode embedded asset: {filename}")
    if capture.shape[:2] != (expected_side, expected_side):
        capture = cv2.resize(capture, (expected_side, expected_side), interpolation=cv2.INTER_AREA)
    return capture


def _gray_float(image: NDArray[Any]) -> NDArray[Any]:
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 and image.shape[2] >= 3 else image
    return gray.astype(np.float32) / 255.0


def _overlaps(candidate: _Candidate, selected: _Candidate) -> bool:
    radius = 0.5 * max(candidate.scale, selected.scale)
    return abs(candidate.x - selected.x) < radius and abs(candidate.y - selected.y) < radius


_TEMPLATE_SCALES = tuple(range(16, 120, 2))


class GeminiEngine:
    """Project-native detector and mask builder for the white Gemini sparkle."""

    _CORE_ALPHA_FRAC = 0.8
    _SPARKLE_FP_CONF = 0.65
    _SPARKLE_FP_MARGIN = 5.0
    _SPARKLE_FP_GRAD = 0.55
    _SPARKLE_KEEP_CONF = 0.52
    _SPARKLE_WHITE_SAT = 0.20
    _CORNER_PROMOTE_NCC = 0.85
    _CORNER_PROMOTE_FRAC = 0.20
    _CORNER_PROMOTE_MIN = 96
    _CORNER_PROMOTE_MAX = 384
    _SELECT_TOPK = 3
    _MASK_ALPHA = 0.04
    _MASK_DILATE_FRAC = 0.18

    def __init__(self, logo_value: float = 255.0) -> None:
        self.logo_value = logo_value
        self._alpha_small = _calculate_alpha_map(_load_capture("gemini_bg_48.png", 48))
        self._alpha_large = _calculate_alpha_map(_load_capture("gemini_bg_96.png", 96))
        self._tmpl_cache: dict[int, NDArray[Any]] = {
            side: cv2.resize(self._alpha_large, (side, side), interpolation=cv2.INTER_AREA) for side in _TEMPLATE_SCALES
        }

    def get_alpha_map(self, size: WatermarkSize) -> NDArray[Any]:
        return self._alpha_small if size is WatermarkSize.SMALL else self._alpha_large

    def get_interpolated_alpha(self, size_px: int) -> NDArray[Any]:
        if size_px == self._alpha_large.shape[1]:
            return self._alpha_large.copy()
        method = cv2.INTER_LINEAR if size_px > self._alpha_large.shape[1] else cv2.INTER_AREA
        return cv2.resize(self._alpha_large, (size_px, size_px), interpolation=method)

    def _scan_scales(self, gray: NDArray[Any]) -> Iterator[tuple[int, float, tuple[int, int]]]:
        """Yield the strongest normalized template match at every usable scale."""
        height, width = gray.shape[:2]
        for side, template in self._tmpl_cache.items():
            if side > height or side > width:
                continue
            response = cv2.matchTemplate(gray, template, cv2.TM_CCOEFF_NORMED)
            _minimum, maximum, _min_location, max_location = cv2.minMaxLoc(response)
            yield side, float(maximum), max_location

    def _global_candidates(self, image: NDArray[Any]) -> list[_Candidate]:
        height, width = image.shape[:2]
        search_side = min(height, width, 512)
        origin_x, origin_y = width - search_side, height - search_side
        gray = _gray_float(image[origin_y:height, origin_x:width])
        ranked = sorted(
            (
                (
                    score * min(1.0, (side / 96.0) ** 0.5),
                    _Candidate(side, origin_x + location[0], origin_y + location[1], score),
                )
                for side, score, location in self._scan_scales(gray)
            ),
            key=lambda item: (item[0], item[1].scale, item[1].spatial, item[1].x, item[1].y),
            reverse=True,
        )
        selected: list[_Candidate] = []
        for _weighted, candidate in ranked:
            if any(_overlaps(candidate, prior) for prior in selected):
                continue
            selected.append(candidate)
            if len(selected) == self._SELECT_TOPK:
                break
        return selected

    def _score_candidate(self, image: NDArray[Any], candidate: _Candidate) -> _Candidate:
        if candidate.spatial < 0.25:
            return candidate
        gradient, variance = self._grad_var_scores(image, candidate.scale, candidate.x, candidate.y)
        return _Candidate(candidate.scale, candidate.x, candidate.y, candidate.spatial, gradient, variance)

    def detect_watermark(
        self,
        image: NDArray[Any],
        force_size: WatermarkSize | None = None,
        *,
        trust_provenance: bool = False,
    ) -> DetectionResult:
        """Return the strongest sparkle-shaped bottom-right candidate."""
        result = DetectionResult()
        if image is None or image.size == 0:
            return result

        source = image_io.to_bgr(image)
        height, width = source.shape[:2]
        result.size = force_size or get_watermark_size(width, height)
        candidates = self._global_candidates(source)
        promoted = self._corner_promote(source, candidates[0].spatial if candidates else -1.0)
        if promoted is not None:
            candidates.append(_Candidate(promoted[0], promoted[1], promoted[2], promoted[3]))
        if not candidates:
            return result

        best = max((self._score_candidate(source, candidate) for candidate in candidates), key=lambda item: item.fused)
        result.region = (best.x, best.y, best.scale, best.scale)
        result.spatial_score = float(best.spatial)
        result.gradient_score = float(best.gradient)
        result.variance_score = float(best.variance)

        confidence = best.fused
        if best.spatial >= 0.25 and confidence < self._SPARKLE_FP_CONF and not trust_provenance:
            confidence = self._apply_false_positive_gate(source, best, confidence)
        result.confidence = float(np.clip(confidence, 0.0, 1.0))
        result.detected = result.confidence >= 0.35
        return result

    def _apply_false_positive_gate(self, image: NDArray[Any], candidate: _Candidate, confidence: float) -> float:
        alpha = self.get_interpolated_alpha(candidate.scale)
        position = (candidate.x, candidate.y)
        margin = self._core_ring_margin(image, alpha, position)
        low_margin = margin is not None and margin < self._SPARKLE_FP_MARGIN
        low_gradient = candidate.gradient < self._SPARKLE_FP_GRAD
        if not low_margin and not low_gradient:
            return confidence
        saturation = self._core_saturation(image, alpha, position)
        neutral_core = not low_margin and saturation is not None and saturation <= self._SPARKLE_WHITE_SAT
        if confidence >= self._SPARKLE_KEEP_CONF and neutral_core:
            return confidence
        logger.debug(
            "Sparkle candidate demoted: confidence=%.3f, margin=%s, gradient=%.3f, saturation=%s",
            confidence,
            margin,
            candidate.gradient,
            saturation,
        )
        return min(confidence, 0.30)

    def _grad_var_scores(self, image: NDArray[Any], scale: int, pos_x: int, pos_y: int) -> tuple[float, float]:
        height, width = image.shape[:2]
        x2, y2 = min(width, pos_x + scale), min(height, pos_y + scale)
        region = image[pos_y:y2, pos_x:x2]
        gray = _gray_float(region)
        alpha = self.get_interpolated_alpha(scale)[: y2 - pos_y, : x2 - pos_x]

        image_edges = cv2.magnitude(
            cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3),
            cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3),
        )
        alpha_edges = cv2.magnitude(
            cv2.Sobel(alpha, cv2.CV_32F, 1, 0, ksize=3),
            cv2.Sobel(alpha, cv2.CV_32F, 0, 1, ksize=3),
        )
        response = cv2.matchTemplate(image_edges, alpha_edges, cv2.TM_CCOEFF_NORMED)
        _minimum, gradient, _min_location, _max_location = cv2.minMaxLoc(response)

        variance = 0.0
        reference_height = min(pos_y, scale)
        if reference_height > 8:
            reference = image[pos_y - reference_height : pos_y, pos_x:x2]
            reference_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY) if reference.ndim == 3 else reference
            _mean, region_std = cv2.meanStdDev((gray * 255.0).astype(np.uint8))
            _reference_mean, reference_std = cv2.meanStdDev(reference_gray)
            if reference_std[0][0] > 5.0:
                variance = float(np.clip(1.0 - region_std[0][0] / reference_std[0][0], 0.0, 1.0))
        return float(gradient), variance

    def _corner_promote(self, image: NDArray[Any], current_raw_ncc: float) -> tuple[int, int, int, float] | None:
        height, width = image.shape[:2]
        desired = round(min(width, height) * self._CORNER_PROMOTE_FRAC)
        side = min(min(width, height), max(self._CORNER_PROMOTE_MIN, min(self._CORNER_PROMOTE_MAX, desired)))
        origin_x, origin_y = width - side, height - side
        matches = self._scan_scales(_gray_float(image[origin_y:height, origin_x:width]))
        best = max(matches, key=lambda item: item[1], default=None)
        if best is None or best[1] < self._CORNER_PROMOTE_NCC or best[1] <= current_raw_ncc:
            return None
        return best[0], origin_x + best[2][0], origin_y + best[2][1], float(best[1])

    def footprint_mask(
        self,
        image: NDArray[Any],
        *,
        force: bool = False,
        dilate: int | None = None,
        region: tuple[int, int, int, int] | None = None,
    ) -> NDArray[Any] | None:
        """Build a full-frame mask from a resolved or newly detected sparkle."""
        if image is None or image.size == 0:
            return None
        source = image_io.to_bgr(image)
        height, width = source.shape[:2]
        if region is not None:
            x, y, scale = region[:3]
        else:
            detection = self.detect_watermark(source)
            if detection.detected:
                x, y, scale = detection.region[:3]
            elif force:
                config = get_watermark_config(width, height)
                x, y = config.get_position(width, height)
                scale = config.logo_size
            else:
                return None

        placed = self._footprint_indices(self.get_interpolated_alpha(scale), (x, y), source.shape)
        if placed is None:
            return None
        alpha, (y1, y2, x1, x2) = placed
        silhouette = (alpha > self._MASK_ALPHA).astype(np.uint8) * 255
        if not silhouette.any():
            return None
        mask = np.zeros((height, width), dtype=np.uint8)
        mask[y1:y2, x1:x2] = silhouette
        radius = dilate if dilate is not None else max(13, int(scale * self._MASK_DILATE_FRAC))
        if radius > 0:
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * radius + 1, 2 * radius + 1))
            mask = cv2.dilate(mask, kernel)
        return mask

    def _footprint_indices(
        self,
        alpha_map: NDArray[Any],
        position: tuple[int, int],
        image_shape: tuple[int, ...],
    ) -> tuple[NDArray[Any], tuple[int, int, int, int]] | None:
        x, y = position
        alpha_height, alpha_width = alpha_map.shape[:2]
        image_height, image_width = image_shape[:2]
        x1, y1 = max(0, x), max(0, y)
        x2, y2 = min(image_width, x + alpha_width), min(image_height, y + alpha_height)
        if x1 >= x2 or y1 >= y2:
            return None
        alpha_x, alpha_y = x1 - x, y1 - y
        clipped = alpha_map[alpha_y : alpha_y + y2 - y1, alpha_x : alpha_x + x2 - x1]
        return clipped, (y1, y2, x1, x2)

    def _core_mask_and_box(
        self,
        image: NDArray[Any],
        alpha_map: NDArray[Any],
        position: tuple[int, int],
    ) -> tuple[NDArray[Any], NDArray[Any], tuple[int, int, int, int], float] | None:
        placed = self._footprint_indices(alpha_map, position, image.shape)
        if placed is None:
            return None
        alpha, bounds = placed
        peak = float(alpha.max())
        if peak < 0.2:
            return None
        core = alpha >= peak * self._CORE_ALPHA_FRAC
        if not core.any():
            return None
        y1, y2, x1, x2 = bounds
        return core, image[y1:y2, x1:x2], bounds, peak

    def _core_and_bg(
        self,
        image: NDArray[Any],
        alpha_map: NDArray[Any],
        position: tuple[int, int],
    ) -> tuple[float, float, float] | None:
        sample = self._core_mask_and_box(image, alpha_map, position)
        if sample is None:
            return None
        core, _box, (y1, y2, x1, x2), peak = sample
        height, width = image.shape[:2]
        padding = int((x2 - x1) * 0.7)
        ry1, ry2 = max(0, y1 - padding), min(height, y2 + padding)
        rx1, rx2 = max(0, x1 - padding), min(width, x2 + padding)
        luminance = image[ry1:ry2, rx1:rx2].astype(np.float32).mean(axis=2)
        fy1, fy2, fx1, fx2 = y1 - ry1, y2 - ry1, x1 - rx1, x2 - rx1
        core_value = float(np.percentile(luminance[fy1:fy2, fx1:fx2][core], 75))
        background = np.ones(luminance.shape, dtype=bool)
        background[fy1:fy2, fx1:fx2] = False
        if background.sum() < 10:
            return None
        return core_value, float(np.median(luminance[background])), peak

    def _core_ring_margin(
        self,
        image: NDArray[Any],
        alpha_map: NDArray[Any],
        position: tuple[int, int],
    ) -> float | None:
        sample = self._core_and_bg(image, alpha_map, position)
        return None if sample is None else sample[0] - sample[1]

    def _core_saturation(
        self,
        image: NDArray[Any],
        alpha_map: NDArray[Any],
        position: tuple[int, int],
    ) -> float | None:
        sample = self._core_mask_and_box(image, alpha_map, position)
        if sample is None:
            return None
        core, box, _bounds, _peak = sample
        pixels = box[core].astype(np.float32)
        brightest = pixels.max(axis=1)
        darkest = pixels.min(axis=1)
        return float(np.median((brightest - darkest) / (brightest + 1.0)))


@functools.lru_cache(maxsize=1)
def _shared_engine() -> GeminiEngine:
    return GeminiEngine()


def detect_sparkle_confidence(image_path: Path, *, image: NDArray[Any] | None = None) -> float | None:
    """Return the local sparkle confidence, or None when decoding fails."""
    decoded = image if image is not None else image_io.imread(image_path)
    if decoded is None:
        return None
    return float(_shared_engine().detect_watermark(decoded).confidence)