"""Render the game's own vehicles into isometric sprites for the planner.

Project Zomboid draws cars as 3D models, so there are no car sprites in the tile packs. This renders the real models
(media/models_X/vehicles/*.fbx, text FBX) with their real paint shells (media/textures/Vehicles/*shell.png) and the
game's wheel mesh (media/models/Vehicles_Wheel.txt), placed and scaled from each vehicle script
(media/scripts/generated/vehicles), at the game's isometric camera: a tile is 128 x 64 pixels at 2x and one tile of
height is 78.4 pixels. Four facings per car: the nose toward S, N, E or W.

Used by extract_game_textures.py; the results go into the same private atlas as the other sprites.
"""
import re
from pathlib import Path

import numpy as np
from PIL import Image

# (id, planner name, vehicle script name, paint colour). The script supplies model, scale, offset, shell and wheels.
VEHICLES = [
    ('sedan', 'Chevalier Nyala (sedan)', 'CarNormal', (122, 30, 28)),
    ('wagon', 'Station wagon', 'CarStationWagon', (58, 82, 112)),
    ('pickup', 'Pick-up truck', 'PickUpTruck', (150, 138, 112)),
    ('suv', 'SUV', 'SUV', (46, 62, 44)),
    ('modern', 'Modern sedan', 'ModernCar02', (196, 196, 190)),
    ('compact', 'Compact car', 'SmallCar02', (170, 120, 40)),
    ('van', 'Van', 'Van', (205, 205, 198)),
    ('luxury', 'Luxury car', 'CarLuxury', (28, 28, 32)),
]
FACES = {'S': (0.0, 1.0), 'N': (0.0, -1.0), 'E': (1.0, 0.0), 'W': (-1.0, 0.0)}
PX, PY, PZ = 64.0, 32.0, 78.4                  # 2x screen pixels per tile along x/y and per tile of height
LIGHT = np.array([-0.45, -0.55, 0.70]); LIGHT /= np.linalg.norm(LIGHT)


def _arr(text, name):
    m = re.search(name + r': \*\d+ \{\s*a: ([^}]*)\}', text)
    return np.array([float(v) for v in m.group(1).replace('\n', '').split(',') if v.strip()]) if m else None


def _fbx_binary(data):
    """Minimal binary FBX reader: returns the first Geometry node's Vertices, PolygonVertexIndex, UV and UVIndex."""
    import struct, zlib
    ver = struct.unpack_from('<I', data, 23)[0]; wide = ver >= 7500
    def node(pos):
        if wide:
            end, nprops, plen = struct.unpack_from('<QQQ', data, pos); pos += 24
        else:
            end, nprops, plen = struct.unpack_from('<III', data, pos); pos += 12
        nl = data[pos]; name = data[pos + 1:pos + 1 + nl].decode('latin1'); pos += 1 + nl
        if end == 0:
            return None, pos
        props = []
        for _ in range(nprops):
            t = chr(data[pos]); pos += 1
            if t in 'fdlib':
                n, enc, clen = struct.unpack_from('<III', data, pos); pos += 12
                raw = data[pos:pos + clen]; pos += clen
                if enc: raw = zlib.decompress(raw)
                props.append(np.frombuffer(raw, {'f': '<f4', 'd': '<f8', 'l': '<i8', 'i': '<i4', 'b': '<u1'}[t]).copy())
            elif t in 'SR':
                n = struct.unpack_from('<I', data, pos)[0]; pos += 4; props.append(data[pos:pos + n]); pos += n
            else:
                fmt = {'Y': '<h', 'C': '<?', 'I': '<i', 'F': '<f', 'D': '<d', 'L': '<q'}[t]
                props.append(struct.unpack_from(fmt, data, pos)[0]); pos += struct.calcsize(fmt)
        kids = []
        while pos < end:
            k, pos = node(pos)
            if k is None: break
            kids.append(k)
        return {'name': name, 'props': props, 'kids': kids}, end
    pos, found = 27, None
    def find(n, name):
        if n['name'] == name: return n
        for k in n['kids']:
            r = find(k, name)
            if r: return r
    while pos < len(data) - 160:
        n, pos = node(pos)
        if n is None: break
        found = find(n, 'Geometry')
        if found: break
    get = lambda n, name: next(k for k in n['kids'] if k['name'] == name)['props'][0]
    uvn = next(k for k in found['kids'] if k['name'] == 'LayerElementUV')
    return get(found, 'Vertices'), get(found, 'PolygonVertexIndex'), get(uvn, 'UV'), get(uvn, 'UVIndex')


def fbx_mesh(path):
    """Vertices, triangles, UVs and UV triangles of the first mesh in an FBX 7.x file (text or binary)."""
    data = open(path, 'rb').read()
    if data.startswith(b'Kaydara FBX Binary'):
        v, pvi, uv, uvi = _fbx_binary(data)
        v = v.reshape(-1, 3); pvi = pvi.astype(int); uv = uv.reshape(-1, 2); uvi = uvi.astype(int)
    else:
        t = data.decode('latin1')
        g = t[t.index('Geometry: '):]
        v = _arr(g, 'Vertices').reshape(-1, 3)
        pvi = _arr(g, 'PolygonVertexIndex').astype(int)
        uvsec = g[g.index('LayerElementUV: 0'):]
        uv = _arr(uvsec, 'UV').reshape(-1, 2); uvi = _arr(uvsec, 'UVIndex').astype(int)
    tris, tuv, poly, puv = [], [], [], []
    for k, i in enumerate(pvi):
        poly.append(~i if i < 0 else i); puv.append(uvi[k])
        if i < 0:
            for j in range(1, len(poly) - 1):
                tris.append((poly[0], poly[j], poly[j + 1])); tuv.append((puv[0], puv[j], puv[j + 1]))
            poly, puv = [], []
    return v, np.array(tris), uv, np.array(tuv)


def pz_mesh(path):
    """The game's own text mesh format (media/models/*.txt): positions, triangles, UVs."""
    L = [l.strip() for l in open(path) if l.strip()]
    i = L.index('# Vertex Count:'); n = int(L[i + 1]); j = L.index('# Vertex Buffer:') + 1
    pos = [[float(x) for x in L[j + 3 * k].split(',')] for k in range(n)]
    uv = [[float(x) for x in L[j + 3 * k + 2].split(',')] for k in range(n)]
    rest = [r for r in L[j + 3 * n:] if not r.startswith('#')]
    nf = int(rest[0]); faces = [[int(x) for x in r.split(',')] for r in rest[1:1 + nf]]
    return np.array(pos), np.array(faces), np.array(uv)


def vehicle_script(scripts, name):
    """Merge a vehicle's script with its template and read what rendering needs."""
    texts = []
    for p in (scripts / 'generated' / 'vehicles').glob('*.txt'):
        t = p.read_text(errors='ignore')
        if re.search(r'vehicle\s+' + name + r'\s*\n\s*\{', t):
            texts.append(t)
    t = '\n'.join(sorted(texts, key=lambda s: 'template vehicle' in s))
    num = lambda s: [float(x) for x in s.split()]
    model = re.search(r'model\s*\{[^}]*?file\s*=\s*(\w+)[^}]*?scale\s*=\s*([\d.]+)[^}]*?offset\s*=\s*([-\d. ]+)', t, re.S)
    shell = re.search(r'skin\s*\{[^}]*?texture\s*=\s*([\w/]+)', t, re.S)
    ext = re.search(r'extents\s*=\s*([-\d. ]+),', t)
    wheels = {m.group(1): num(m.group(2)) for m in re.finditer(r'wheel\s+(\w+)\s*\{[^}]*?offset\s*=\s*([-\d. ]+),', t, re.S)}
    radius = re.search(r'wheel\s+\w+\s*\{[^}]*?radius\s*=\s*([\d.]+)', t, re.S)
    return {'file': model.group(1), 'scale': float(model.group(2)), 'offset': num(model.group(3)), 'shell': shell.group(1),
            'extents': num(ext.group(1)), 'wheels': wheels, 'radius': float(radius.group(1)) if radius else .15}


def rasterise(tris, size, anchor, ss=2):
    """Z-buffered textured triangles -> RGBA image. tris: list of (screen xy[3], depth[3], uv[3], texture, paint, normal)."""
    W, H = size[0] * ss, size[1] * ss
    rgb = np.zeros((H, W, 3), np.float32); alpha = np.zeros((H, W), np.float32); zb = np.full((H, W), -1e9, np.float32)
    for P, dep, uv, tex, paint, n in tris:
        P = (P + anchor) * ss
        x0, y0 = np.floor(P.min(0)).astype(int); x1, y1 = np.ceil(P.max(0)).astype(int)
        x0, y0, x1, y1 = max(x0, 0), max(y0, 0), min(x1, W - 1), min(y1, H - 1)
        if x1 < x0 or y1 < y0:
            continue
        xs, ys = np.meshgrid(np.arange(x0, x1 + 1) + .5, np.arange(y0, y1 + 1) + .5)
        (ax, ay), (bx, by), (cx, cy) = P
        den = (by - cy) * (ax - cx) + (cx - bx) * (ay - cy)
        if abs(den) < 1e-9:
            continue
        w0 = ((by - cy) * (xs - cx) + (cx - bx) * (ys - cy)) / den
        w1 = ((cy - ay) * (xs - cx) + (ax - cx) * (ys - cy)) / den
        w2 = 1 - w0 - w1
        inside = (w0 >= -1e-4) & (w1 >= -1e-4) & (w2 >= -1e-4)
        if not inside.any():
            continue
        d = w0 * dep[0] + w1 * dep[1] + w2 * dep[2]
        sub = zb[y0:y1 + 1, x0:x1 + 1]
        m = inside & (d > sub)
        if not m.any():
            continue
        u = w0 * uv[0][0] + w1 * uv[1][0] + w2 * uv[2][0]; v = w0 * uv[0][1] + w1 * uv[1][1] + w2 * uv[2][1]
        th, tw = tex.shape[:2]
        tx = np.clip((u % 1.0) * tw, 0, tw - 1).astype(int); ty = np.clip((1 - v % 1.0) * th, 0, th - 1).astype(int)
        texel = tex[ty, tx]
        if isinstance(paint, str):                      # 'cut': alpha-tested texture (clothing layers, masked body)
            m = m & (texel[..., 3] > 127)
            if not m.any():
                continue
        col = texel[..., :3] * (texel[..., 3:4] / 255.0) + np.array(paint, np.float32) * (1 - texel[..., 3:4] / 255.0) if paint is not None and not isinstance(paint, str) else texel[..., :3]
        shade = .58 + .42 * max(0.0, float(np.dot(n, LIGHT)))
        rgb[y0:y1 + 1, x0:x1 + 1][m] = col[m] * shade
        alpha[y0:y1 + 1, x0:x1 + 1][m] = 255 if paint is not None else np.maximum(texel[..., 3], 0)[m] if False else 255
        sub[m] = d[m]
    img = Image.fromarray(np.dstack([np.clip(rgb, 0, 255), alpha]).astype(np.uint8), 'RGBA')
    return img.resize(size, Image.LANCZOS)


def arch_top(W, H, L, tris, ox, oz, above):
    """Lowest point of the body skin over a wheel: the section of the mesh at the wheel's axle, outboard of its centre."""
    P = np.column_stack([W, H, L])
    a, b = P[tris].reshape(-1, 3), P[np.roll(tris, -1, axis=1)].reshape(-1, 3)
    cross = ((a[:, 2] - oz) * (b[:, 2] - oz) <= 0) & (a[:, 2] != b[:, 2])
    a, b = a[cross], b[cross]
    p = a + ((oz - a[:, 2]) / (b[:, 2] - a[:, 2]))[:, None] * (b - a)
    p = p[(p[:, 0] * np.sign(ox) > abs(ox)) & (p[:, 1] > above)]
    return float(p[:, 1].min()) if len(p) else None


def render_vehicle(media, script, paint, face):
    base = media / 'models_X' / 'vehicles' / (script['file'] + '.fbx')
    v, tris, uv, tuv = fbx_mesh(base)
    shell = np.asarray(Image.open(media / 'textures' / (script['shell'] + '.png')).convert('RGBA')).astype(np.float32)
    wheel_tex = np.asarray(Image.open(media / 'textures' / 'Vehicles' / 'vehicle_wheel.png').convert('RGBA')).astype(np.float32)
    wp, wf, wuv = pz_mesh(media / 'models' / 'Vehicles_Wheel.txt')
    wuv = np.column_stack([wuv[:, 0], 1 - wuv[:, 1]])        # the game's text meshes store texture v top-down
    s = script['scale']; ext = script['extents']; off = script['offset']
    k = ext[2] / (v[:, 1].max() - v[:, 1].min())            # mesh units -> script units, from the vehicle's length
    rad = max(abs(wp[:, 1]).max(), script['radius'])
    W, H, L = v[:, 0] * k, v[:, 2] * k + off[1], v[:, 1] * k
    # The script's wheel points are suspension mounts at full extension, so the wheels hang up to a radius below the
    # arches. Seat each wheel just under the top of its arch in the body mesh, as the game shows a car at rest.
    seat = {}
    for name, (ox, oy, oz) in script['wheels'].items():
        top = arch_top(W, H, L, tris, ox, oz, off[1] + oy)
        seat[name] = off[1] + oy if top is None else max(off[1] + oy, top - rad - .012)
    ground = (min(seat.values()) if seat else off[1] - .3) - rad
    F = np.array(FACES[face]); Lft = np.array([F[1], -F[0]])
    def world(width, height, length):                         # script units (left, up, forward) -> tiles (x, y, h)
        xy = np.outer(length * s, F) + np.outer(width * s, Lft)
        return np.column_stack([xy, (height - ground) * s])
    body = world(v[:, 0] * k, v[:, 2] * k + off[1], v[:, 1] * k)
    parts = [(body, tris, uv, tuv, shell, paint)]
    for name, (ox, oy, oz) in script['wheels'].items():
        side = 1 if ox > 0 else -1
        wpos = world(wp[:, 0] * side + ox, wp[:, 1] + seat[name], wp[:, 2] + oz)
        parts.append((wpos, wf, wuv, wf, wheel_tex, None))
    out, pts = [], []
    for P3, T, UV, TUV, tex, pnt in parts:
        sx = (P3[:, 0] - P3[:, 1]) * PX; sy = (P3[:, 0] + P3[:, 1]) * PY - P3[:, 2] * PZ
        dep = .612 * P3[:, 0] + .612 * P3[:, 1] + .5 * P3[:, 2]
        S2 = np.column_stack([sx, sy]); pts.append(S2)
        for a, b in zip(T, TUV):
            e1, e2 = P3[a[1]] - P3[a[0]], P3[a[2]] - P3[a[0]]
            n = np.cross(e1, e2); ln = np.linalg.norm(n)
            n = n / ln if ln > 1e-12 else np.array([0, 0, 1.0])
            if np.dot(n, [.612, .612, .5]) < 0:
                n = -n
            out.append((S2[a], dep[a], UV[b], tex, pnt, n))
    allp = np.vstack(pts)
    mn, mx = np.floor(allp.min(0)) - 4, np.ceil(allp.max(0)) + 4
    size = (int(mx[0] - mn[0]), int(mx[1] - mn[1]))
    img = rasterise(out, size, -mn)
    return img, -mn[0], -mn[1]                                 # anchor: where the car's centre on the ground is


def render_all(media):
    """{sprite name: (image, ox, oy)} with offsets in the atlas's 128 x 256 frame convention, plus the vehicle list."""
    scripts = media / 'scripts'
    imgs, cars = {}, []
    for vid, name, sname, paint in VEHICLES:
        try:
            sc = vehicle_script(scripts, sname)
        except (AttributeError, ValueError):
            continue
        faces = {}
        for f in FACES:
            img, ax, ay = render_vehicle(media, sc, paint, f)
            n = f'vehicle_{vid}_{f}'
            imgs[n] = (img, int(round(64 - ax)), int(round(192 - ay)))
            faces[f] = n
        cars.append({'id': vid, 'name': name, 'script': sname, 'faces': faces,
                     'length': round(sc['extents'][2] * sc['scale'], 2), 'width': round(sc['extents'][0] * sc['scale'], 2)})
    return imgs, cars


if __name__ == '__main__':
    import sys
    media = Path(sys.argv[1]) if len(sys.argv) > 1 and sys.argv[1] else Path.home() / 'Library/Application Support/Steam/steamapps/common/ProjectZomboid/Project Zomboid.app/Contents/Java/media'
    imgs, cars = render_all(media)
    sheet = Image.new('RGBA', (1600, 1000), (40, 44, 40, 255)); x = y = 10; rowh = 0
    for n, (im, ox, oy) in imgs.items():
        if x + im.width > 1600: x, y, rowh = 10, y + rowh + 10, 0
        sheet.alpha_composite(im, (x, y)); x += im.width + 10; rowh = max(rowh, im.height)
    out = Path(sys.argv[2]) if len(sys.argv) > 2 else Path('vehicles-check.png')
    sheet.save(out); print(len(imgs), 'vehicle sprites', [c['id'] for c in cars])