"""
PBR Tangent-Space Normal Map Synthesis Utility for M4 Assets.
Provides vectorized Sobel filter computation with guaranteed Blue channel mean > 128.0.
"""

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


def compute_tangent_normal_map(albedo_img: Image.Image, strength: float = 2.2) -> Image.Image:
    """Computes tangent-space normal map using 3x3 Sobel operator with B mean > 128 guarantee."""
    gray = albedo_img.convert("L")
    arr = np.array(gray, dtype=np.float32) / 255.0

    padded = np.pad(arr, 1, mode="edge")
    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:]
    ) * strength

    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:]
    ) * 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)

    if albedo_img.mode == "RGBA":
        alpha = np.array(albedo_img.split()[-1])
        # In transparent areas, set flat tangent normal (128, 128, 255)
        r[alpha == 0] = 128
        g[alpha == 0] = 128
        b[alpha == 0] = 255
        return Image.merge("RGBA", (Image.fromarray(r), Image.fromarray(g), Image.fromarray(b), Image.fromarray(alpha)))

    return Image.merge("RGB", (Image.fromarray(r), Image.fromarray(g), Image.fromarray(b)))