"""Signature cleaning: white-background removal and alpha-mask extraction.

Public handwritten signature datasets are typically flat scans: dark ink
strokes on an off-white/paper background, with no alpha channel. This module
converts such a scan into a clean RGBA "stamp" -- background fully
transparent, ink strokes opaque-to-translucent according to how dark the
original pixel was (preserving natural stroke texture: thin/faint pen
lift-offs stay translucent, heavy strokes stay opaque) -- ready to be
transformed and composited onto a document by :class:`~src.placement.placement_engine.PlacementEngine`
and :class:`~src.rendering.renderer.Renderer`.
"""

from __future__ import annotations

import cv2
import numpy as np

from src.utils.logging_utils import get_logger

logger = get_logger(__name__)

DEFAULT_INK_COLOR: tuple[int, int, int] = (18, 18, 24)


class SignaturePreprocessor:
    """Cleans a raw signature scan into a tightly-cropped, transparent RGBA stamp.

    Example:
        >>> pre = SignaturePreprocessor()
        >>> rgba = pre.process(raw_signature_array)
    """

    def __init__(
        self,
        ink_color: tuple[int, int, int] = DEFAULT_INK_COLOR,
        min_component_area: int = 4,
        pad_px: int = 6,
        alpha_gamma: float = 0.85,
    ) -> None:
        self.ink_color = ink_color
        self.min_component_area = min_component_area
        self.pad_px = pad_px
        self.alpha_gamma = alpha_gamma

    def process(self, raw_image: np.ndarray, ink_color: "tuple[int, int, int] | None" = None) -> np.ndarray:
        """Convert a raw signature scan into a clean, tightly-cropped RGBA stamp.

        Args:
            raw_image: The raw signature image as loaded (grayscale ``(H, W)``,
                RGB ``(H, W, 3)``, or RGBA ``(H, W, 4)`` uint8).
            ink_color: Optional override for the flat ink color used for the
                RGB channels (alpha still carries the natural stroke texture).

        Returns:
            An ``(H', W', 4)`` uint8 RGBA array, cropped tightly to the ink
            content with a small transparent pad.
        """
        gray, existing_alpha = self._to_gray(raw_image)
        gray = cv2.medianBlur(gray, 3)
        alpha = self._compute_alpha(gray)
        if existing_alpha is not None:
            alpha = np.minimum(alpha, existing_alpha)
        alpha = self._remove_speckles(alpha)
        rgba = self._compose_rgba(alpha, ink_color or self.ink_color)
        rgba = self._crop_to_content(rgba, pad=self.pad_px)
        return rgba

    @staticmethod
    def _to_gray(raw_image: np.ndarray) -> tuple[np.ndarray, "np.ndarray | None"]:
        if raw_image.ndim == 2:
            return raw_image.copy(), None
        if raw_image.shape[2] == 4:
            rgb = raw_image[:, :, :3]
            alpha = raw_image[:, :, 3]
            gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
            return gray, alpha
        gray = cv2.cvtColor(raw_image[:, :, :3], cv2.COLOR_RGB2GRAY)
        return gray, None

    @staticmethod
    def _compute_alpha(gray: np.ndarray) -> np.ndarray:
        gray_f = gray.astype(np.float64)
        background = float(np.percentile(gray_f, 92))
        dark_pixels = gray_f[gray_f < background - 8]
        dark_level = float(np.percentile(dark_pixels, 5)) if dark_pixels.size > 20 else max(background - 140, 0.0)
        denom = max(background - dark_level, 1e-3)

        raw = (background - gray_f) / denom
        raw = np.clip(raw, 0.0, 1.0)
        # Discard near-background noise (sub-threshold jitter from scanning artifacts).
        raw[raw < 0.06] = 0.0
        softened = raw**0.85
        return (softened * 255.0).astype(np.uint8)

    def _remove_speckles(self, alpha: np.ndarray) -> np.ndarray:
        binary = (alpha > 15).astype(np.uint8)
        n_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
        if n_labels <= 1:
            return alpha
        cleaned = alpha.copy()
        for label in range(1, n_labels):
            area = stats[label, cv2.CC_STAT_AREA]
            if area < self.min_component_area:
                cleaned[labels == label] = 0
        return cleaned

    @staticmethod
    def _compose_rgba(alpha: np.ndarray, ink_color: tuple[int, int, int]) -> np.ndarray:
        h, w = alpha.shape
        rgb = np.empty((h, w, 3), dtype=np.uint8)
        rgb[:, :] = ink_color
        return np.dstack([rgb, alpha])

    @staticmethod
    def _crop_to_content(rgba: np.ndarray, pad: int) -> np.ndarray:
        alpha = rgba[:, :, 3]
        ys, xs = np.where(alpha > 4)
        if ys.size == 0:
            logger.warning("Signature preprocessing found no ink pixels; returning a 1x1 transparent stamp.")
            return np.zeros((1, 1, 4), dtype=np.uint8)
        y1, y2 = max(ys.min() - pad, 0), min(ys.max() + pad + 1, rgba.shape[0])
        x1, x2 = max(xs.min() - pad, 0), min(xs.max() + pad + 1, rgba.shape[1])
        return rgba[y1:y2, x1:x2].copy()