#!/usr/bin/env python3
"""Background removal using rembg (ML-based, robust to non-uniform backgrounds).

This module provides a single, high-quality implementation backed by the
open-source `rembg` library. It outputs PNG with alpha transparency.

Usage in code:
    from api_sdk.ml_background_remover import remove_background_ml
    remove_background_ml("input.jpg", "output.png", alpha_matting=True)

Notes:
    - Requires `pip install rembg`.
    - Does not perform any network calls.
"""

from pathlib import Path
from typing import Optional

from PIL import Image
import numpy as np

try:
    import cv2  # optional but recommended for mask ops
except Exception:  # pragma: no cover
    cv2 = None


def remove_background_ml(
    image_path: str,
    output_path: Optional[str] = None,
    *,
    alpha_matting: bool = False,
    alpha_matting_foreground_threshold: int = 240,
    alpha_matting_background_threshold: int = 10,
    alpha_matting_erosion_size: int = 10,
    model: Optional[str] = None,
    # Guided preservation options
    keep_mask_path: Optional[str] = None,
    keep_mask_dilate: int = 0,
    alpha_support_threshold: int = 40,
    nonwhite_keep_threshold: Optional[int] = None,
    auto_keep_from_nonwhite: bool = False,
    auto_seed_erosion: int = 2,
) -> Path:
    """Remove background using rembg and save as PNG with transparency.

    Args:
        image_path: Input image path.
        output_path: Output path; defaults to '<stem>_transparent.png'.
        alpha_matting: Enable alpha matting refinement for cleaner edges.
        alpha_matting_foreground_threshold: Foreground threshold (0-255).
        alpha_matting_background_threshold: Background threshold (0-255).
        alpha_matting_erosion_size: Erosion size for trimap refinement.
        model: Optional rembg model name (e.g., 'isnet-general-use').

    Returns:
        Path to the saved PNG file with transparent background.
    """
    try:
        from rembg import remove  # type: ignore
    except Exception as e:
        raise RuntimeError(
            "rembg is required for ML background removal. Install with 'pip install rembg'."
        ) from e

    # Load input image
    img = Image.open(image_path)

    # rembg accepts PIL Image and returns PIL Image (in common versions)
    result = remove(
        img,
        alpha_matting=alpha_matting,
        alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,
        alpha_matting_background_threshold=alpha_matting_background_threshold,
        alpha_matting_erosion_size=alpha_matting_erosion_size,
        model=model,
    )

    # Determine output path
    if output_path is None:
        base = Path(image_path).stem
        output_path = f"{base}_transparent.png"

    out_path = Path(output_path)

    # Normalize to RGBA image
    if not isinstance(result, Image.Image):
        # Assume bytes-like
        with open(out_path, "wb") as f:
            f.write(result)
        # Reopen for possible guided fusion below
        rgba = Image.open(out_path).convert("RGBA")
    else:
        rgba = result.convert("RGBA")

    # Optional: guided preservation
    if any([keep_mask_path, auto_keep_from_nonwhite, nonwhite_keep_threshold is not None]):
        if cv2 is None:
            raise RuntimeError("OpenCV (cv2) is required for guided preservation. Install opencv-python.")

        src_rgb = Image.open(image_path).convert("RGB")
        src_arr = np.array(src_rgb)
        out_arr = np.array(rgba)
        alpha = out_arr[:, :, 3].astype(np.uint8)

        # Support mask: where model already sees plausible foreground
        support = (alpha >= int(alpha_support_threshold)).astype(np.uint8) * 255

        # Optional: treat clearly non-white as support as well
        if nonwhite_keep_threshold is not None:
            diff = np.abs(src_arr.astype(np.int16) - 255)
            maxdiff = diff.max(axis=2).astype(np.uint8)
            _, nonwhite = cv2.threshold(maxdiff, int(nonwhite_keep_threshold), 255, cv2.THRESH_BINARY)
            support = cv2.max(support, nonwhite)

        # Seeds from keep mask or auto non-white erosion
        seeds = None
        if keep_mask_path:
            m = Image.open(keep_mask_path).convert("L")
            seeds = (np.array(m) > 0).astype(np.uint8) * 255
        elif auto_keep_from_nonwhite:
            diff = np.abs(src_arr.astype(np.int16) - 255)
            maxdiff = diff.max(axis=2).astype(np.uint8)
            _, seeds = cv2.threshold(maxdiff, 25, 255, cv2.THRESH_BINARY)
            k = np.ones((3, 3), np.uint8)
            if auto_seed_erosion > 0:
                seeds = cv2.erode(seeds, k, iterations=int(auto_seed_erosion))

        if seeds is not None:
            # Optional seed dilation (to be less sparse)
            if keep_mask_dilate > 0:
                k = np.ones((3, 3), np.uint8)
                seeds = cv2.dilate(seeds, k, iterations=int(keep_mask_dilate))

            # Morphological reconstruction (geodesic dilation) under support
            # Initialize marker as seeds ∧ support
            marker = cv2.bitwise_and(seeds, support)
            kernel = np.ones((3, 3), np.uint8)
            while True:
                dil = cv2.dilate(marker, kernel, iterations=1)
                new_marker = cv2.bitwise_and(dil, support)
                if np.array_equal(new_marker, marker):
                    break
                marker = new_marker

            # Force preserve reconstructed region
            out_arr[:, :, 3] = np.maximum(out_arr[:, :, 3], marker)
            rgba = Image.fromarray(out_arr, "RGBA")

    # Save final output
    rgba.save(out_path, "PNG")

    return out_path
