"""Mesh loading, UV-position-map rasterization and texture baking helpers.

Re-implements Paint3D's `UV_pos_render` (which relies on kaolin) with plain numpy:
every texel inside a UV triangle gets the interpolated 3D position of that surface
point, with the mesh centred, scaled into a sphere of radius `shape_scale` (0.6 in
Paint3D's config), and mapped from [-1, 1] to [0, 1] RGB. Empty texels are black.
"""
import numpy as np
import trimesh
from PIL import Image
from scipy import ndimage

SHAPE_SCALE = 0.6


def load_mesh(path: str):
    """Load a mesh file and return (scene, geometry_name, mesh) for the largest UV-mapped geometry.

    Meshes without UVs are unwrapped with xatlas.
    """
    scene = trimesh.load(path, force="scene", process=False)
    if not scene.geometry:
        raise ValueError("No geometry found in the uploaded file.")
    name, mesh = max(scene.geometry.items(), key=lambda kv: len(kv[1].faces))
    uv = getattr(mesh.visual, "uv", None)
    if uv is None or len(uv) != len(mesh.vertices):
        import xatlas

        vmap, faces, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
        mesh = trimesh.Trimesh(
            vertices=mesh.vertices[vmap], faces=faces,
            visual=trimesh.visual.TextureVisuals(uv=uvs), process=False,
        )
        scene.geometry[name] = mesh
    return scene, name, mesh


def normalized_vertices(mesh: trimesh.Trimesh) -> np.ndarray:
    v = np.asarray(mesh.vertices, dtype=np.float64)
    v = v - v.mean(axis=0)
    v = v / np.linalg.norm(v, axis=1).max()
    return v * SHAPE_SCALE


def rasterize_uv_position(mesh: trimesh.Trimesh, res: int = 1024):
    """Return (uv_pos HxWx3 float32 in [0,1], coverage mask HxW bool)."""
    uv = np.asarray(mesh.visual.uv, dtype=np.float64) % 1.0000001
    pos = normalized_vertices(mesh)
    faces = np.asarray(mesh.faces)

    # texel space: x = u*res, y = (1-v)*res  (image row 0 == v=1, OBJ/glTF texture convention)
    tri_uv = np.stack([uv[faces, 0] * res, (1.0 - uv[faces, 1]) * res], axis=-1)  # F,3,2
    tri_pos = pos[faces]  # F,3,3

    out = np.zeros((res, res, 3), dtype=np.float32)
    zbuf = np.zeros((res, res), dtype=bool)

    lo = np.clip(np.floor(tri_uv.min(axis=1) - 0.5), 0, res - 1).astype(np.int64)  # F,2
    hi = np.clip(np.ceil(tri_uv.max(axis=1) + 0.5), 0, res - 1).astype(np.int64)
    size = hi - lo + 1
    extent = size.max(axis=1)

    # batch triangles with similar bbox sizes so each batch is a dense (B, S, S) grid
    order = np.argsort(extent)
    for bucket in np.array_split(order, max(1, len(order) // 4096)):
        if len(bucket) == 0:
            continue
        S = int(extent[bucket].max())
        if S * S * len(bucket) > 64_000_000:  # huge triangles: fall back to smaller chunks
            for sub in np.array_split(bucket, int(np.ceil(S * S * len(bucket) / 64_000_000))):
                _raster_batch(sub, S, tri_uv, tri_pos, lo, res, out, zbuf)
        else:
            _raster_batch(bucket, S, tri_uv, tri_pos, lo, res, out, zbuf)

    out = np.clip(out / 2 + 0.5, 0, 1)
    out[~zbuf] = 0
    return out, zbuf


def _raster_batch(idx, S, tri_uv, tri_pos, lo, res, out, cov):
    grid = np.arange(S)
    px = lo[idx, 0][:, None, None] + grid[None, None, :] + 0.5  # B,1,S
    py = lo[idx, 1][:, None, None] + grid[None, :, None] + 0.5  # B,S,1
    a, b, c = (tri_uv[idx, k][:, None, None, :] for k in range(3))
    v0, v1 = b - a, c - a
    den = v0[..., 0] * v1[..., 1] - v1[..., 0] * v0[..., 1]
    den = np.where(np.abs(den) < 1e-12, 1e-12, den)
    dx, dy = px - a[..., 0], py - a[..., 1]
    w1 = (dx * v1[..., 1] - v1[..., 0] * dy) / den
    w2 = (v0[..., 0] * dy - dx * v0[..., 1]) / den
    w0 = 1 - w1 - w2
    eps = -1e-4
    inside = (w0 >= eps) & (w1 >= eps) & (w2 >= eps)
    X = np.broadcast_to(px - 0.5, inside.shape).astype(np.int64)
    Y = np.broadcast_to(py - 0.5, inside.shape).astype(np.int64)
    inside &= (X < res) & (Y < res)
    b_i, _, _ = np.nonzero(inside)
    w = np.stack([w0[inside], w1[inside], w2[inside]], axis=-1)
    p = np.einsum("nk,nkc->nc", w, tri_pos[idx][b_i])
    out[Y[inside], X[inside]] = p
    cov[Y[inside], X[inside]] = True


def dilate_texture(img: np.ndarray, mask: np.ndarray, pad: int = 16) -> np.ndarray:
    """Fill empty texels near UV islands with the nearest valid colour (avoids seams when mip-mapping)."""
    if mask.all():
        return img
    dist, (iy, ix) = ndimage.distance_transform_edt(~mask, return_indices=True)
    filled = img[iy, ix]
    keep = mask | (dist <= pad)
    res = img.copy()
    res[keep] = filled[keep]
    res[~keep] = filled[~keep]  # fill the rest too; harmless and cleaner in viewers
    return res


def export_textured(scene, name, mesh, texture: Image.Image, keep_normal: bool, out_path: str):
    old = getattr(mesh.visual, "material", None)
    normal = getattr(old, "normalTexture", None) if keep_normal else None
    if normal is not None and max(normal.size) > 2048:
        normal = normal.resize((2048, 2048), Image.LANCZOS)
    mat = trimesh.visual.material.PBRMaterial(
        baseColorTexture=texture, normalTexture=normal,
        metallicFactor=0.0, roughnessFactor=0.85, doubleSided=True,
    )
    mesh.visual = trimesh.visual.TextureVisuals(uv=mesh.visual.uv, material=mat)
    scene.geometry[name] = mesh
    scene.export(out_path)
    return out_path