"""
FreeExile Alpha Mask Segmentation & Transparency Cleaner.
Removes opaque background remnants using OpenCV GrabCut and edge-feathering,
producing 100% clean transparent PNG sprites and normal maps.
Eliminates all white halo / milky blur issues during Canvas animations.
"""

from __future__ import annotations

import os
from pathlib import Path
import cv2
import numpy as np
from PIL import Image

WORKSPACE_ROOT = Path(__file__).resolve().parent.parent.parent
CLIENT_ASSETS = WORKSPACE_ROOT / "client" / "webapp" / "assets"
MONSTERS_DIR = CLIENT_ASSETS / "monsters"
CHARACTERS_DIR = CLIENT_ASSETS / "characters"

TARGET_SPRITES = [
    MONSTERS_DIR / "mob_skeleton_warrior.png",
    MONSTERS_DIR / "mob_feral_hellhound.png",
    MONSTERS_DIR / "mob_rot_crawler.png",
    MONSTERS_DIR / "mob_shadow_wraith.png",
    MONSTERS_DIR / "mob_primal_cannibal.png",
    MONSTERS_DIR / "mob_ironhide_behemoth.png",
    MONSTERS_DIR / "mob_corrupted_raptor.png",
    MONSTERS_DIR / "mob_bramble_treant.png",
    MONSTERS_DIR / "mob_tomb_lord.png",
    MONSTERS_DIR / "mob_flesh_abomination.png",
    MONSTERS_DIR / "boss_blood_bone_ravager.png",
    MONSTERS_DIR / "boss_abyssal_tyrant.png",
    CHARACTERS_DIR / "savage_primal_exile.png",
]


def compute_tangent_normal_map_cv(rgba: np.ndarray, strength: float = 3.0) -> np.ndarray:
    """Computes tangent-space normal map from RGBA image."""
    b, g, r, a = cv2.split(rgba)
    gray = cv2.cvtColor(cv2.merge([b, g, r]), cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0

    # Sobel gradients
    dx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3) * strength
    dy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3) * strength
    dz = np.ones_like(gray)

    len_vec = np.sqrt(dx**2 + dy**2 + dz**2)
    len_vec[len_vec == 0] = 1.0

    nx = dx / len_vec
    ny = dy / len_vec
    nz = dz / len_vec

    norm_r = ((nx * 0.5 + 0.5) * 255).astype(np.uint8)
    norm_g = ((ny * 0.5 + 0.5) * 255).astype(np.uint8)
    norm_b = ((nz * 0.5 + 0.5) * 255).astype(np.uint8)

    return cv2.merge([norm_b, norm_g, norm_r, a])


def extract_clean_transparent_silhouette(img_path: Path) -> None:
    """Applies GrabCut foreground extraction and feathers boundary."""
    if not img_path.exists():
        print(f"[!] File not found: {img_path}")
        return

    print(f"[*] Processing transparency for: {img_path.name}...")
    img = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
    if img is None:
        return

    rgb = img[:, :, :3]
    h, w = rgb.shape[:2]

    # Initialize mask and background/foreground models
    mask = np.zeros((h, w), np.uint8)
    bgd_model = np.zeros((1, 65), np.float64)
    fgd_model = np.zeros((1, 65), np.float64)

    # 1. Bounding box focusing on the central character/monster
    margin_x = int(w * 0.08)
    margin_y = int(h * 0.06)
    rect = (margin_x, margin_y, w - 2 * margin_x, h - 2 * margin_y)

    # 2. Run GrabCut iterations
    cv2.grabCut(rgb, mask, rect, bgd_model, fgd_model, 3, cv2.GC_INIT_WITH_RECT)

    # 3. Create foreground binary mask (Probable FG + Definite FG)
    fg_mask = np.where((mask == 1) | (mask == 3), 255, 0).astype(np.uint8)

    # Clean morphological noise (fill small holes in torso, remove tiny stray specs)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE, kernel)

    # Smooth feathered edge
    fg_smooth = cv2.GaussianBlur(fg_mask, (5, 5), 0)

    # 4. Merge RGB with cleaned alpha
    b, g, r = cv2.split(rgb)
    clean_rgba = cv2.merge([b, g, r, fg_smooth])

    # Save cleaned PNG
    cv2.imwrite(str(img_path), clean_rgba)

    # Recompute and save matching normal map
    normal_path = img_path.parent / f"{img_path.stem}_normal.png"
    clean_normal = compute_tangent_normal_map_cv(clean_rgba, strength=3.0)
    cv2.imwrite(str(normal_path), clean_normal)

    fg_ratio = fg_smooth.mean() / 255.0
    print(f" -> Completed: {img_path.name} (Foreground ratio: {fg_ratio:.1%}, normal map updated).")


def run_transparency_cleanup() -> None:
    print("=================================================================")
    print(" FREEEXILE: TRANSPARENT SPRITE SEGMENTATION & CLEANUP ENGINE    ")
    print("=================================================================")
    for path in TARGET_SPRITES:
        extract_clean_transparent_silhouette(path)
    print("=================================================================")
    print(" [SUCCESS] ALL SPRITES CONVERTED TO TRUE TRANSPARENT SILHOUETTES ")
    print("=================================================================")


if __name__ == "__main__":
    run_transparency_cleanup()