from __future__ import annotations

import numpy as np

from .chroma import BackgroundModel
from .deps import require_cv2
from .settings import DespillSettings


def despill(
    rgb: np.ndarray,
    alpha: np.ndarray,
    model: BackgroundModel,
    settings: DespillSettings,
) -> tuple[np.ndarray, np.ndarray]:
    """Remove background-colour spill from the foreground.

    The background's Lab a*/b* vector gives a chroma direction; foreground chroma
    projecting along it beyond a margin is treated as spill and pulled back toward
    neutral, keeping luminance so hair texture and shading survive. This works for
    any flat background: a saturated primary (green/blue screen) reduces to classic
    channel suppression, while low-chroma pastels -- where no single RGB channel
    dominates -- are handled by the same projection.

    Returns the corrected RGB (uint8) and the per-pixel despill weight (float32).
    """
    if not settings.enabled:
        return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)

    bg_lab = np.asarray(model.lab_center, dtype=np.float32)
    va, vb = float(bg_lab[1]), float(bg_lab[2])
    mag = float(np.hypot(va, vb))
    if mag < 1e-3:  # near-neutral background: no chroma direction, nothing to despill
        return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)
    da, db = va / mag, vb / mag

    cv2 = require_cv2()
    rgb_f = rgb.astype(np.float32) / 255.0
    foreground = alpha > settings.edge_low
    if settings.edge_expand_radius > 0:
        radius = settings.edge_expand_radius
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1))
        foreground = cv2.dilate(foreground.astype(np.uint8), kernel).astype(bool)

    lab = cv2.cvtColor(rgb_f, cv2.COLOR_RGB2LAB)
    A, B = lab[..., 1], lab[..., 2]
    proj = A * da + B * db                       # chroma toward the background hue
    excess = np.maximum(proj - settings.color_margin_lab, 0.0)

    # Full strength across the opaque range; ease off only where mostly transparent.
    alpha_weight = np.maximum(np.clip(alpha * 2.0, 0.0, 1.0), settings.alpha_weight_floor)
    color_weight = np.clip(excess / max(settings.color_margin_lab * 8.0, 1e-6), 0.0, 1.0)
    mask = foreground.astype(np.float32) * alpha_weight * color_weight
    reduce = excess * mask * settings.strength   # pull chroma back along the bg axis

    lab_out = lab.copy()
    lab_out[..., 1] = A - reduce * da
    lab_out[..., 2] = B - reduce * db
    out = cv2.cvtColor(lab_out, cv2.COLOR_LAB2RGB)
    return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32)