from __future__ import annotations

import numpy as np

from .chroma import BackgroundModel
from .settings import DespillSettings


def despill_green(
    rgb: np.ndarray,
    alpha: np.ndarray,
    bg_confidence: np.ndarray,
    model: BackgroundModel,
    settings: DespillSettings,
) -> tuple[np.ndarray, np.ndarray]:
    rgb_f = rgb.astype(np.float32) / 255.0
    if not settings.enabled:
        return rgb.copy(), np.zeros(alpha.shape, dtype=np.float32)

    edge = (alpha > settings.edge_low) & (alpha < settings.edge_high)
    r = rgb_f[..., 0]
    g = rgb_f[..., 1]
    b = rgb_f[..., 2]
    neutral_green = np.maximum(r, b) + settings.green_excess_margin
    excess = np.maximum(g - neutral_green, 0.0)
    edge_weight = np.clip((1.0 - np.abs(alpha - 0.5) * 2.0), 0.0, 1.0)
    bg_weight = np.clip(bg_confidence, 0.0, 1.0)
    mask = edge.astype(np.float32) * edge_weight * np.maximum(bg_weight, 0.25)

    out = rgb_f.copy()
    out[..., 1] = g - excess * mask * settings.strength

    # Very light compensation toward the non-green channels to avoid gray fringes.
    bg_green = float(model.rgb_center[1])
    compensation = excess * mask * settings.strength * min(0.25, bg_green * 0.15)
    out[..., 0] = np.clip(out[..., 0] + compensation * 0.5, 0.0, 1.0)
    out[..., 2] = np.clip(out[..., 2] + compensation * 0.5, 0.0, 1.0)

    return np.clip(out * 255.0, 0, 255).astype(np.uint8), mask.astype(np.float32)