"""
FreeExile Town Environment Assets Processor.
Converts AI-generated isometric town structures on white backgrounds into
high-fidelity transparent PNGs and Tangent-Space 3D Normal Maps.
"""

import sys
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"

ASSET_MAP = [
    ("ancient_golden_tree_1790931719454.jpg", "tree_golden_ancient.png"),
    ("blacksmith_forge_1790931745649.jpg", "forge_blacksmith.png"),
    ("market_stall_outcast_1790931768957.jpg", "stall_market_outcast.png"),
    ("wuxia_tiled_house_1790931789770.jpg", "house_traditional_tile.png"),
    ("wuxia_props_cluster_1790931814236.jpg", "props_town_cluster.png"),
]


def compute_tangent_normal_map(rgba: np.ndarray, strength: float = 2.5) -> 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

    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_white_background(img_bgr: np.ndarray, thresh: int = 245) -> np.ndarray:
    """Cleans solid white background with flood fill from 4 corners and edge feathering."""
    h, w = img_bgr.shape[:2]
    # Check whiteness
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    is_white = gray >= thresh

    # Flood fill from the 4 corners to only remove external background, not white inside props
    mask = np.zeros((h + 2, w + 2), np.uint8)
    flood_img = img_bgr.copy()
    diff = (12, 12, 12)
    corners = [(0, 0), (w - 1, 0), (0, h - 1), (w - 1, h - 1), (w // 2, 0), (w // 2, h - 1), (0, h // 2), (w - 1, h // 2)]
    for pt in corners:
        if gray[pt[1], pt[0]] >= 220:
            cv2.floodFill(flood_img, mask, pt, (0, 0, 0), diff, diff, flags=8 | (255 << 8) | cv2.FLOODFILL_MASK_ONLY)

    bg_mask = mask[1:h+1, 1:w+1]
    # Invert to get foreground mask
    fg_mask = cv2.bitwise_not(bg_mask)

    # Clean stray specs
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    fg_mask = cv2.morphologyEx(fg_mask, cv2.MORPH_CLOSE, kernel)
    fg_mask = cv2.GaussianBlur(fg_mask, (3, 3), 0)

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


def process_all():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    for src_name, out_name in ASSET_MAP:
        src_path = BRAIN_DIR / src_name
        if not src_path.exists():
            print(f"[!] Source not found: {src_path}")
            continue

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

        rgba = extract_white_background(img)
        normal = compute_tangent_normal_map(rgba)

        out_rgba_path = OUTPUT_DIR / out_name
        out_normal_path = OUTPUT_DIR / out_name.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}")
        print(f"  [+] Saved Normal: {out_normal_path}")


if __name__ == "__main__":
    process_all()