#!/usr/bin/env python3
"""Create clean transparent game-asset cutouts with BiRefNet.

This tool is intended for Novamailio-generated WhaleTown assets that come back
on a plain matte background. It uses BiRefNet for the alpha mask, then fills
transparent/semitransparent edge RGB from confident foreground pixels so Godot
texture filtering cannot reveal the original matte color.
"""

from __future__ import annotations

import argparse
from collections import deque
import os
from pathlib import Path
import sys
from typing import Iterable, Tuple

import numpy as np
from PIL import Image, ImageDraw


DEFAULT_MODEL = "ZhengPeng7/BiRefNet"
DEFAULT_HF_ENDPOINT = "https://hf-mirror.com"
DEFAULT_SIZE = 1024


def _die(message: str, code: int = 1) -> None:
    print(f"Error: {message}", file=sys.stderr)
    raise SystemExit(code)


def _warn(message: str) -> None:
    print(f"Warning: {message}", file=sys.stderr)


def _import_ml_deps() -> Tuple[object, object, object]:
    try:
        import torch
        from torchvision import transforms
        from transformers import AutoModelForImageSegmentation
    except ImportError as exc:
        _die(
            "Missing BiRefNet dependencies. Install them with:\n"
            "  python3 -m pip install torch torchvision transformers timm einops kornia scipy\n"
            f"Original import error: {exc}"
        )
    return torch, transforms, AutoModelForImageSegmentation


def _select_device(torch: object, requested: str) -> str:
    if requested != "auto":
        return requested
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def _load_birefnet(model_name: str, device: str, torch: object, auto_model: object) -> object:
    model = auto_model.from_pretrained(model_name, trust_remote_code=True)
    model.to(device)
    model.eval()
    return model


def _run_birefnet(
    image: Image.Image,
    *,
    model_name: str,
    device: str,
    input_size: int,
) -> Image.Image:
    torch, transforms, auto_model = _import_ml_deps()
    selected_device = _select_device(torch, device)
    print(f"BiRefNet device: {selected_device}", file=sys.stderr)
    print(f"BiRefNet model: {model_name}", file=sys.stderr)
    model = _load_birefnet(model_name, selected_device, torch, auto_model)

    transform = transforms.Compose(
        [
            transforms.Resize(
                (input_size, input_size),
                interpolation=transforms.InterpolationMode.BILINEAR,
            ),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    input_tensor = transform(image).unsqueeze(0).to(selected_device)
    with torch.no_grad():
        output = model(input_tensor)
        if isinstance(output, (list, tuple)):
            prediction = output[-1]
        elif hasattr(output, "logits"):
            prediction = output.logits
        else:
            prediction = output
        prediction = prediction.sigmoid().detach().float().cpu()[0]
        if prediction.ndim == 3:
            prediction = prediction.squeeze(0)
    mask = transforms.ToPILImage()(prediction)
    return mask.resize(image.size, Image.Resampling.LANCZOS).convert("L")


def _stabilize_mask(
    mask: Image.Image,
    *,
    low_cut: float,
    high_span: float,
    hard_low: float,
    hard_high: float,
) -> Image.Image:
    mask_arr = np.asarray(mask).astype(np.float32) / 255.0
    mask_arr = np.clip((mask_arr - low_cut) / high_span, 0.0, 1.0)
    mask_arr = np.where(mask_arr > hard_high, 1.0, mask_arr)
    mask_arr = np.where(mask_arr < hard_low, 0.0, mask_arr)
    return Image.fromarray((mask_arr * 255).astype(np.uint8))


def _connected_matte_mask(image: Image.Image, threshold: float) -> np.ndarray:
    """Fallback mask for comparing or rescuing BiRefNet failures."""
    rgb = np.asarray(image.convert("RGB"))
    arr = rgb.astype(np.int32)
    h, w = arr.shape[:2]
    strips = np.concatenate(
        [
            arr[:80, :, :].reshape(-1, 3),
            arr[max(0, h - 80) : h, :, :].reshape(-1, 3),
            arr[:, :80, :].reshape(-1, 3),
            arr[:, max(0, w - 80) : w, :].reshape(-1, 3),
        ],
        axis=0,
    )
    background = np.median(strips, axis=0).astype(np.int32)
    color_dist = np.sqrt(((arr - background) ** 2).sum(axis=2))
    candidate = color_dist <= threshold
    visited = np.zeros((h, w), dtype=bool)
    queue: deque[Tuple[int, int]] = deque()
    for x in range(w):
        for y in (0, h - 1):
            if candidate[y, x] and not visited[y, x]:
                visited[y, x] = True
                queue.append((y, x))
    for y in range(h):
        for x in (0, w - 1):
            if candidate[y, x] and not visited[y, x]:
                visited[y, x] = True
                queue.append((y, x))
    while queue:
        y, x = queue.popleft()
        for ny, nx in ((y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)):
            if 0 <= ny < h and 0 <= nx < w and (not visited[ny, nx]) and candidate[ny, nx]:
                visited[ny, nx] = True
                queue.append((ny, nx))
    return ~visited


def _decontaminate_edge_rgb(rgb: np.ndarray, alpha: np.ndarray, confidence: int) -> np.ndarray:
    foreground = alpha > confidence
    if not foreground.any():
        _warn("Mask has no confident foreground pixels; edge decontamination skipped.")
        return rgb
    try:
        from scipy import ndimage
    except ImportError:
        _warn("scipy is missing; install scipy for edge RGB decontamination.")
        return rgb

    _, indices = ndimage.distance_transform_edt(~foreground, return_indices=True)
    nearest_rgb = rgb[indices[0], indices[1]]
    alpha_f = alpha.astype(np.float32) / 255.0
    edge_mix = np.clip((0.98 - alpha_f) / 0.98, 0.0, 1.0)[..., None]
    replace_strength = np.where(alpha_f[..., None] < 0.98, edge_mix, 0.0)
    cleaned = rgb * (1.0 - replace_strength) + nearest_rgb * replace_strength
    return np.clip(cleaned, 0, 255).astype(np.uint8)


def _compose_cutout(image: Image.Image, mask: Image.Image, confidence: int) -> Image.Image:
    rgb = np.asarray(image.convert("RGB")).astype(np.float32)
    alpha = np.asarray(mask).astype(np.uint8)
    cleaned_rgb = _decontaminate_edge_rgb(rgb, alpha, confidence)
    return Image.fromarray(np.dstack([cleaned_rgb, alpha]))


def _make_preview(cutout: Image.Image, output: Path, scale: float) -> None:
    w, h = cutout.size
    small = cutout.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
    canvas = Image.new("RGB", (small.width * 2 + 48, small.height + 48), (238, 238, 238))

    magenta_plate = Image.new("RGB", small.size, (198, 76, 190))
    magenta_plate.paste(small, (0, 0), small)

    checker = Image.new("RGB", small.size, (230, 230, 230))
    draw = ImageDraw.Draw(checker)
    step = 32
    for y in range(0, small.height, step):
        for x in range(0, small.width, step):
            if ((x // step) + (y // step)) % 2 == 0:
                draw.rectangle([x, y, x + step - 1, y + step - 1], fill=(190, 206, 224))
    checker.paste(small, (0, 0), small)

    canvas.paste(magenta_plate, (16, 32))
    canvas.paste(checker, (small.width + 32, 32))
    output.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(output)


def _magenta_stats(image: Image.Image) -> Tuple[int, int, int]:
    arr = np.asarray(image.convert("RGBA")).astype(np.int16)
    rgb = arr[:, :, :3]
    alpha = arr[:, :, 3]
    semi = (alpha > 0) & (alpha < 255)
    magenta = (alpha > 0) & (rgb[:, :, 0] > 120) & (rgb[:, :, 1] < 60) & (rgb[:, :, 2] > 90)
    semi_magenta = semi & magenta
    return int(semi.sum()), int(magenta.sum()), int(semi_magenta.sum())


def parse_args(argv: Iterable[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", required=True, help="Source RGB/RGBA image with matte background.")
    parser.add_argument("--output", required=True, help="Transparent PNG output path.")
    parser.add_argument("--mask-out", help="Optional alpha mask output path.")
    parser.add_argument("--preview-out", help="Optional preview sheet output path.")
    parser.add_argument("--model", default=DEFAULT_MODEL, help=f"HF model id. Default: {DEFAULT_MODEL}")
    parser.add_argument("--hf-endpoint", default=DEFAULT_HF_ENDPOINT, help="Hugging Face endpoint/mirror.")
    parser.add_argument("--device", choices=["auto", "cpu", "mps", "cuda"], default="auto")
    parser.add_argument("--input-size", type=int, default=DEFAULT_SIZE, help="Square BiRefNet input size.")
    parser.add_argument("--low-cut", type=float, default=0.025, help="Low alpha normalization cut.")
    parser.add_argument("--high-span", type=float, default=0.94, help="Alpha normalization span.")
    parser.add_argument("--hard-low", type=float, default=0.015, help="Values below this become transparent.")
    parser.add_argument("--hard-high", type=float, default=0.985, help="Values above this become opaque.")
    parser.add_argument("--edge-confidence", type=int, default=245, help="Confident foreground alpha for RGB fill.")
    parser.add_argument("--preview-scale", type=float, default=0.52)
    parser.add_argument("--fallback-connected-matte", action="store_true", help="Use simple connected matte mask instead of BiRefNet.")
    parser.add_argument("--fallback-threshold", type=float, default=28.0)
    return parser.parse_args(list(argv))


def main(argv: Iterable[str]) -> int:
    args = parse_args(argv)
    input_path = Path(args.input)
    output_path = Path(args.output)
    if not input_path.exists():
        _die(f"Input image not found: {input_path}")

    os.environ.setdefault("HF_ENDPOINT", args.hf_endpoint)
    image = Image.open(input_path).convert("RGB")
    if args.fallback_connected_matte:
        foreground = _connected_matte_mask(image, args.fallback_threshold)
        mask = Image.fromarray(foreground.astype(np.uint8) * 255)
    else:
        mask = _run_birefnet(
            image,
            model_name=args.model,
            device=args.device,
            input_size=args.input_size,
        )
        mask = _stabilize_mask(
            mask,
            low_cut=args.low_cut,
            high_span=args.high_span,
            hard_low=args.hard_low,
            hard_high=args.hard_high,
        )

    cutout = _compose_cutout(image, mask, args.edge_confidence)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    cutout.save(output_path)

    if args.mask_out:
        mask_path = Path(args.mask_out)
        mask_path.parent.mkdir(parents=True, exist_ok=True)
        mask.save(mask_path)
    if args.preview_out:
        _make_preview(cutout, Path(args.preview_out), args.preview_scale)

    semi, magenta, semi_magenta = _magenta_stats(cutout)
    print(f"wrote {output_path}")
    if args.mask_out:
        print(f"wrote {args.mask_out}")
    if args.preview_out:
        print(f"wrote {args.preview_out}")
    print(f"alpha_bbox={cutout.getchannel('A').getbbox()}")
    print(f"semi_transparent_pixels={semi}")
    print(f"visible_magenta_like_pixels={magenta}")
    print(f"semi_magenta_like_pixels={semi_magenta}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))