"""Software rasterizer for the exported mosswart mesh.

Loads the JSON dumped by tools/MosswartArt, textures dumped by tools/IconExtract,
and renders a lit, shaded, supersampled image. Deliberately simple: z-buffer,
perspective camera, Lambert key/fill/rim + Blinn specular for wet skin.
"""
import json, glob, os, math
import numpy as np
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", ".."))
TEXDIR = os.path.join(ROOT, "tools", "IconExtract", "bin", "Release", "net10.0", "out")


# ---------------------------------------------------------------- data loading
def load_textures(ids):
    out = {}
    for tid in ids:
        hits = glob.glob(os.path.join(TEXDIR, tid + "_*.png"))
        if not hits:
            continue
        im = Image.open(hits[0]).convert("RGBA")
        out[tid] = np.asarray(im, dtype=np.float32) / 255.0
    return out


def load_mesh(path, parts_filter=None, ignore_transform=False):
    doc = json.load(open(path))
    P, N, UV, TRI, TEX = [], [], [], [], []
    for part in doc["parts"]:
        if parts_filter is not None and part["index"] not in parts_filter:
            continue
        m = np.array(part["m"], dtype=np.float64).reshape(4, 4)
        if ignore_transform:
            m = np.eye(4)
        # System.Numerics Matrix4x4 is row-vector convention: v' = v * M
        rot = m[:3, :3]
        trans = m[3, :3]
        for sub in part["sub"]:
            v = np.array(sub["v"], dtype=np.float64)
            if len(v) == 0:
                continue
            base = len(P)
            pos = v[:, 0:3] @ rot + trans
            nrm = v[:, 3:6] @ rot
            ln = np.linalg.norm(nrm, axis=1, keepdims=True)
            nrm = np.divide(nrm, np.where(ln == 0, 1, ln))
            P.append(pos)
            N.append(nrm)
            UV.append(v[:, 6:8])
            idx = np.array(sub["i"], dtype=np.int64).reshape(-1, 3) + base
            TRI.append(idx)
            TEX += [sub["tex"]] * len(idx)
    return (np.concatenate(P), np.concatenate(N), np.concatenate(UV),
            np.concatenate(TRI), TEX)


# ---------------------------------------------------------------- math helpers
def look_at(eye, target, up=np.array([0.0, 0.0, 1.0])):
    f = target - eye
    f = f / np.linalg.norm(f)
    s = np.cross(f, up)
    if np.linalg.norm(s) < 1e-9:
        s = np.cross(f, np.array([0.0, 1.0, 0.0]))
    s = s / np.linalg.norm(s)
    u = np.cross(s, f)
    return np.stack([s, u, -f])  # rows: right, up, back


def orbit_eye(target, radius, az_deg, el_deg):
    az, el = math.radians(az_deg), math.radians(el_deg)
    return target + radius * np.array([
        math.cos(el) * math.cos(az),
        math.cos(el) * math.sin(az),
        math.sin(el)])


# ---------------------------------------------------------------- rasterizer
def render(mesh, textures, size=512, ss=3, az=90.0, el=8.0, fov=26.0,
           target=None, radius=None, fit=1.0,
           key=(-0.45, -0.85, 0.45), rim=(0.7, 0.75, -0.15),
           bg=(0, 0, 0, 0), key_col=(1.06, 1.04, 0.92), amb_top=(0.30, 0.34, 0.24),
           amb_bot=(0.10, 0.10, 0.06), rim_col=(0.55, 0.72, 0.30), spec_pow=26.0,
           spec_amt=0.42, key_cam=None, rim_cam=None, fill_cam=None,
           fill_col=(0.0, 0.0, 0.0), exposure=1.0, gamma=1.0, matcaps=None):
    P, N, UV, TRI, TEX = mesh
    W = size * ss

    if target is None:
        lo, hi = P.min(axis=0), P.max(axis=0)
        target = (lo + hi) / 2.0
    if radius is None:
        ext = np.linalg.norm(P.max(axis=0) - P.min(axis=0))
        radius = ext / (2.0 * math.tan(math.radians(fov) / 2.0)) * fit

    eye = orbit_eye(target, radius, az, el)
    V = look_at(eye, target)

    cam = (P - eye) @ V.T
    ncam = N @ V.T

    f = 1.0 / math.tan(math.radians(fov) / 2.0)
    z = -cam[:, 2]
    z = np.where(z < 1e-6, 1e-6, z)
    sx = (cam[:, 0] * f / z * 0.5 + 0.5) * W
    sy = (1.0 - (cam[:, 1] * f / z * 0.5 + 0.5)) * W

    color = np.zeros((W, W, 3), dtype=np.float32)
    alpha = np.zeros((W, W), dtype=np.float32)
    depth = np.full((W, W), 1e30, dtype=np.float32)

    keyd = np.array(key, dtype=np.float64); keyd /= np.linalg.norm(keyd)
    rimd = np.array(rim, dtype=np.float64); rimd /= np.linalg.norm(rimd)
    keyc = V @ keyd
    rimc = V @ rimd
    # Camera-relative light directions make art direction repeatable across
    # camera angles -- the key stays on the same side of the face regardless
    # of where the model happens to be facing in world space.
    if key_cam is not None:
        keyc = np.array(key_cam, dtype=np.float64); keyc /= np.linalg.norm(keyc)
    if rim_cam is not None:
        rimc = np.array(rim_cam, dtype=np.float64); rimc /= np.linalg.norm(rimc)
    fillc = None
    if fill_cam is not None:
        fillc = np.array(fill_cam, dtype=np.float64); fillc /= np.linalg.norm(fillc)
    viewd = np.array([0.0, 0.0, 1.0])
    half = keyc + viewd
    half = half / np.linalg.norm(half)

    for t in range(len(TRI)):
        i0, i1, i2 = TRI[t]
        x0, y0 = sx[i0], sy[i0]
        x1, y1 = sx[i1], sy[i1]
        x2, y2 = sx[i2], sy[i2]
        area = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
        if abs(area) < 1e-9:
            continue
        minx = max(int(math.floor(min(x0, x1, x2))), 0)
        maxx = min(int(math.ceil(max(x0, x1, x2))), W - 1)
        miny = max(int(math.floor(min(y0, y1, y2))), 0)
        maxy = min(int(math.ceil(max(y0, y1, y2))), W - 1)
        if minx > maxx or miny > maxy:
            continue

        xs = np.arange(minx, maxx + 1)
        ys = np.arange(miny, maxy + 1)
        gx, gy = np.meshgrid(xs + 0.5, ys + 0.5)
        w0 = ((x1 - gx) * (y2 - gy) - (x2 - gx) * (y1 - gy)) / area
        w1 = ((x2 - gx) * (y0 - gy) - (x0 - gx) * (y2 - gy)) / area
        w2 = 1.0 - w0 - w1
        inside = (w0 >= 0) & (w1 >= 0) & (w2 >= 0)
        if not inside.any():
            continue

        zt = w0 * z[i0] + w1 * z[i1] + w2 * z[i2]
        sub = depth[miny:maxy + 1, minx:maxx + 1]
        closer = inside & (zt < sub)
        if not closer.any():
            continue

        u = w0 * UV[i0, 0] + w1 * UV[i1, 0] + w2 * UV[i2, 0]
        v = w0 * UV[i0, 1] + w1 * UV[i1, 1] + w2 * UV[i2, 1]
        nx = w0 * ncam[i0, 0] + w1 * ncam[i1, 0] + w2 * ncam[i2, 0]
        ny = w0 * ncam[i0, 1] + w1 * ncam[i1, 1] + w2 * ncam[i2, 1]
        nz = w0 * ncam[i0, 2] + w1 * ncam[i1, 2] + w2 * ncam[i2, 2]
        nl = np.sqrt(nx * nx + ny * ny + nz * nz)
        nl = np.where(nl == 0, 1, nl)
        nx, ny, nz = nx / nl, ny / nl, nz / nl
        # two-sided: AC meshes are frequently single-sided but authored either way
        flip = nz < 0
        nx = np.where(flip, -nx, nx); ny = np.where(flip, -ny, ny); nz = np.where(flip, -nz, nz)

        tex = textures.get(TEX[t])
        if tex is None:
            rgb = np.ones(u.shape + (3,), dtype=np.float32) * 0.6
            ta = np.ones(u.shape, dtype=np.float32)
        else:
            th, tw = tex.shape[0], tex.shape[1]
            fu = (u % 1.0) * (tw - 1)
            fv = (v % 1.0) * (th - 1)
            u0 = np.clip(np.floor(fu).astype(np.int32), 0, tw - 1)
            v0 = np.clip(np.floor(fv).astype(np.int32), 0, th - 1)
            u1 = np.clip(u0 + 1, 0, tw - 1)
            v1 = np.clip(v0 + 1, 0, th - 1)
            du = (fu - u0)[..., None]
            dv = (fv - v0)[..., None]
            texel = (tex[v0, u0] * (1 - du) * (1 - dv) + tex[v0, u1] * du * (1 - dv)
                     + tex[v1, u0] * (1 - du) * dv + tex[v1, u1] * du * dv)
            rgb = texel[..., :3]
            ta = texel[..., 3]

        ndl = np.clip(nx * keyc[0] + ny * keyc[1] + nz * keyc[2], 0, 1)
        ndr = np.clip(nx * rimc[0] + ny * rimc[1] + nz * rimc[2], 0, 1) ** 2.4
        up = np.clip(nz * 0.35 + 0.65, 0, 1)
        amb = (np.array(amb_bot)[None, None, :] +
               (np.array(amb_top) - np.array(amb_bot))[None, None, :] * up[..., None])
        ndh = np.clip(nx * half[0] + ny * half[1] + nz * half[2], 0, 1)
        spec = (ndh ** spec_pow) * spec_amt * ndl

        lightsum = amb + np.array(key_col)[None, None, :] * ndl[..., None]
        if fillc is not None:
            ndf = np.clip(nx * fillc[0] + ny * fillc[1] + nz * fillc[2], 0, 1)
            lightsum = lightsum + np.array(fill_col)[None, None, :] * ndf[..., None]
        lit = (rgb * lightsum
               + np.array(rim_col)[None, None, :] * ndr[..., None]
               + spec[..., None]) * exposure
        if gamma != 1.0:
            lit = np.clip(lit, 0, None) ** gamma

        # Metal: a Lambert term cannot produce chrome, because chrome is almost
        # entirely reflection. Sample a matcap by the camera-space normal instead.
        mc = matcaps.get(TEX[t]) if matcaps else None
        if mc is not None:
            mh, mw = mc.shape[0], mc.shape[1]
            mu = np.clip((nx * 0.5 + 0.5) * (mw - 1), 0, mw - 1).astype(np.int32)
            mv = np.clip((1.0 - (ny * 0.5 + 0.5)) * (mh - 1), 0, mh - 1).astype(np.int32)
            lit = mc[mv, mu][..., :3] * exposure
            ta = np.ones(u.shape, dtype=np.float32)

        m = closer & (ta > 0.35)
        if not m.any():
            continue
        ys_i, xs_i = np.nonzero(m)
        color[miny + ys_i, minx + xs_i] = lit[ys_i, xs_i]
        alpha[miny + ys_i, minx + xs_i] = 1.0
        depth[miny + ys_i, minx + xs_i] = zt[ys_i, xs_i]

    rgba = np.concatenate([np.clip(color, 0, 1), alpha[..., None]], axis=2)
    img = Image.fromarray((rgba * 255).astype(np.uint8), "RGBA")
    return img.resize((size, size), Image.LANCZOS)


def all_texture_ids(path):
    doc = json.load(open(path))
    return sorted({s["tex"] for p in doc["parts"] for s in p["sub"]})