"""
FreeExile New Town Environment Props Processor.
Extracts AI-generated props using morphological contour bounding and hole-filling,
preserving 100% of dark and light internal materials, crops to bounding boxes,
generates PBR normal maps, and exports clean transparent PNGs.
"""

from pathlib import Path
import cv2
import numpy as np

BRAIN_DIR = Path(r"C:\Users\Admin\.gemini\antigravity\brain\629e29bf-683e-4041-b730-afb3479ebffc")
WORKSPACE_ROOT = Path(__file__).resolve().parent.parent.parent
OUTPUT_DIR = WORKSPACE_ROOT / "client" / "webapp" / "assets" / "environment"

PROPS = [
    ("campfire_bonfire_embers_1790934339544.jpg", "campfire_bonfire_embers.png", "black"),
    ("lantern_stone_pagoda_1790934365046.jpg", "lantern_stone_pagoda.png", "black"),
    ("weapon_rack_arsenal_1790934393110.jpg", "weapon_rack_arsenal.png", "black"),
    ("totem_war_banner_1790934420773.jpg", "totem_war_banner.png", "black"),
    ("well_stone_ancient_1790934453051.jpg", "well_stone_ancient.png", "white"),
]


def compute_tangent_normal_map(rgba: np.ndarray, strength: float = 2.5) -> np.ndarray:
    b, g, r, a = cv2.split(rgba)
    gray = cv2.cvtColor(cv2.merge([b, g, r]), cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0

    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_background(img_bgr: np.ndarray, bg_type: str = "black") -> np.ndarray:
    h, w = img_bgr.shape[:2]

    if bg_type == "black":
        val = img_bgr.max(axis=2)
        binary = (val > 14).astype(np.uint8) * 255
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        dilated = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = [c for c in contours if cv2.contourArea(c) > 400]
        mask = np.zeros_like(val, dtype=np.uint8)
        cv2.drawContours(mask, contours, -1, 255, -1)
        fg_mask = cv2.GaussianBlur(mask, (3, 3), 0)
    else:
        gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
        mask_buf = np.zeros((h + 2, w + 2), np.uint8)
        flood = img_bgr.copy()
        corners = [(0, 0), (w - 1, 0), (0, h - 1), (w - 1, h - 1)]
        for pt in corners:
            cv2.floodFill(flood, mask_buf, pt, (0, 0, 0), (12, 12, 12), (12, 12, 12),
                          flags=8 | (255 << 8) | cv2.FLOODFILL_MASK_ONLY)
        bg = mask_buf[1:h+1, 1:w+1]
        fg = (bg == 0).astype(np.uint8) * 255
        contours, _ = cv2.findContours(fg, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        contours = [c for c in contours if cv2.contourArea(c) > 400]
        mask = np.zeros_like(gray, dtype=np.uint8)
        cv2.drawContours(mask, contours, -1, 255, -1)
        fg_mask = cv2.GaussianBlur(mask, (3, 3), 0)

    b, g, r = cv2.split(img_bgr)
    rgba = cv2.merge([b, g, r, fg_mask])

    # Crop to non-transparent bounding box with small margin
    coords = cv2.findNonZero(fg_mask)
    if coords is not None:
        x, y, bw, bh = cv2.boundingRect(coords)
        pad = 6
        x0 = max(0, x - pad)
        y0 = max(0, y - pad)
        x1 = min(w, x + bw + pad)
        y1 = min(h, y + bh + pad)
        rgba = rgba[y0:y1, x0:x1]

    return rgba


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    for src_file, out_file, bg_type in PROPS:
        src_path = BRAIN_DIR / src_file
        if not src_path.exists():
            print(f"[!] File not found: {src_path}")
            continue

        print(f"[*] Processing {src_file} ({bg_type}) -> {out_file}...")
        img = cv2.imread(str(src_path))
        if img is None:
            print(f"[!] Failed to read: {src_path}")
            continue

        rgba = extract_background(img, bg_type)
        normal = compute_tangent_normal_map(rgba)

        out_rgba_path = OUTPUT_DIR / out_file
        out_normal_path = OUTPUT_DIR / out_file.replace(".png", "_normal.png")

        cv2.imwrite(str(out_rgba_path), rgba)
        cv2.imwrite(str(out_normal_path), normal)
        print(f"  [+] Saved RGBA: {out_rgba_path} ({rgba.shape[1]}x{rgba.shape[0]})")
        print(f"  [+] Saved Normal: {out_normal_path}")


if __name__ == "__main__":
    main()