"""
Tangent-space Normal Map Utilities for FreeExile Asset Pipeline.
Provides vectorized Sobel filter operations to synthesize PBR normal maps from albedo/bump.
"""

from __future__ import annotations
import numpy as np
from PIL import Image


def generate_normal_map(albedo_img: Image.Image, strength: float = 2.5) -> Image.Image:
    """Generates tangent-space normal map using vectorized 3x3 Sobel operator."""
    gray = albedo_img.convert("L")
    arr = np.array(gray, dtype=np.float32) / 255.0

    padded = np.pad(arr, 1, mode="edge")
    # Vectorized Sobel convolution
    # X filter: [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]
    dx = (
        -1.0 * padded[0:-2, 0:-2] + 1.0 * padded[0:-2, 2:]
        - 2.0 * padded[1:-1, 0:-2] + 2.0 * padded[1:-1, 2:]
        - 1.0 * padded[2:, 0:-2] + 1.0 * padded[2:, 2:]
    )
    # Y filter: [[-1, -2, -1], [0, 0, 0], [1, 2, 1]]
    dy = (
        -1.0 * padded[0:-2, 0:-2] - 2.0 * padded[0:-2, 1:-1] - 1.0 * padded[0:-2, 2:]
        + 1.0 * padded[2:, 0:-2] + 2.0 * padded[2:, 1:-1] + 1.0 * padded[2:, 2:]
    )

    dx *= strength
    dy *= strength
    dz = np.ones_like(arr)

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

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

    r = ((nx * 0.5 + 0.5) * 255).astype(np.uint8)
    g = ((ny * 0.5 + 0.5) * 255).astype(np.uint8)
    b = ((nz * 0.5 + 0.5) * 255).astype(np.uint8)

    alpha = albedo_img.split()[-1]
    return Image.merge("RGBA", (Image.fromarray(r), Image.fromarray(g), Image.fromarray(b), alpha))