#!/usr/bin/env python3
"""Geometry library for level WADs of "Ein Koenigreich fuer ein Lama" (Croc 2 engine, PC).

Format: see notes/geometry.md. In short:
  TRAK  = tile models (world geometry + collision), u32 n, n x 0x84 header, then data
  MAP   = tile grid: placement (position/rotation/model) + vertex colours (+ strats etc.)
  STPC  = object models (same model format, header 0x8c) + animations + strat scripts
  TEXT  = texture pages + texture table (UV rectangles)

All parsers are reimplemented from groove.exe (FUN_005563f0 TRAK, FUN_0041f770 STPC model,
FUN_0042ac50 MAP, FUN_00406d30 texture table, FUN_00556510 tile renderer).
"""
import math
import struct

# ---------------------------------------------------------------- OBJ file space
# OBJ files (geom_export/geom_import/geom_demo, the editor's Import mesh and OBJ/glTF export) are
# right-handed, y up, in game space: file (x, y, z) = export (x, y, -z). The game shows export space
# left-handed, so this is how Blender shows the level the same way round as the game. Converting
# mirrors, so the corner order reverses (first corner kept: export (0,1,2) <-> file (0,2,1)).


def to_file(p):
    """Export-space point -> OBJ file point."""
    return (p[0], p[1], -p[2])


def file_order(seq):
    """Corner order export <-> file (an involution)."""
    seq = list(seq)
    return seq[:1] + seq[:0:-1]


# ---------------------------------------------------------------- WAD


def wad_read(path):
    d = open(path, 'rb').read()
    total = struct.unpack_from('<I', d, 0)[0]
    assert total == len(d) - 4, 'WAD length field does not match'
    chunks, o = [], 4
    while o + 8 <= len(d):
        tag = d[o:o + 4][::-1].decode('latin1')
        ln = struct.unpack_from('<I', d, o + 4)[0]
        chunks.append([tag, d[o + 8:o + 8 + ln]])
        o += 8 + ln
        if tag == 'END ':
            break
    return chunks


def wad_write(path, chunks):
    body = b''.join(t.encode('latin1')[::-1] + struct.pack('<I', len(c)) + c for t, c in chunks)
    open(path, 'wb').write(struct.pack('<I', len(body)) + body)


def wad_get(chunks, tag):
    for t, c in chunks:
        if t == tag.ljust(4):
            return c
    return None


def wad_set(chunks, tag, data):
    for c in chunks:
        if c[0] == tag.ljust(4):
            c[1] = data
            return
    raise KeyError(tag)


def wfpc_flags(chunks):
    c = wad_get(chunks, 'WFPC')
    return struct.unpack('<I', c)[0] if c else 0


# ---------------------------------------------------------------- models
# Header (TRAK 0x84 bytes, STPC 0x8c bytes):
#   +0x00 3f   centre of the bounding box
#   +0x0c 24f  8 corners of the bounding box (enlarged by 1/128)
#   +0x6c u16  nVerts      +0x6e u16 nFaces
#   +0x70 u32  (verts pointer, 0 in file)   +0x74 u32 (faces pointer, 0)
#   +0x78 u16  nColA  +0x7a u16 nColB  +0x7c u16 nColC   (collision polygons, 32 bytes)
#   +0x7e u16  0      +0x80 u32 (collision pointer, 0)
#   STPC only: +0x84 u16 ?(usually = nVerts or 0)  +0x86 u16 nX (n*2 bytes of data before the verts)
#             +0x88 u32 (pointer, 0)
# Data: [STPC: nX*u16] verts nV*24 (pos 3f, normal 3f), faces nF*28, collision nCol*32
# Face (28): u16 flags, u16 v0,v1,v2, u16 tex, u16 0, 3f plane normal, f d (n.v = d)
# Collision (32): u8 type (bit 0: quad, otherwise triangle), u8 material, 5 x {i8 nx,ny,nz, u8 0,
#   i16 d}: plane 0 = surface, planes 1..4 = edges (pointing inwards). Normals *127, d *4096,
#   coordinates in game space (z mirrored relative to the vertex data!).

CORNERS = [(-1, -1, 1), (1, -1, 1), (-1, 1, 1), (1, 1, 1), (-1, -1, -1), (1, -1, -1), (-1, 1, -1), (1, 1, -1)]


class Model:
    def __init__(self):
        self.verts = []      # [(x,y,z,nx,ny,nz)]
        self.faces = []      # [dict(flags,v=(a,b,c),tex,pad,plane=(nx,ny,nz,d))]
        self.coll = [[], [], []]   # three groups; 32 raw bytes per entry (bytes)
        self.hdr = None      # original header (bytes) for comparison
        self.extra = b''     # STPC: nX*u16
        self.x84 = 0
        self.dirty = False   # set True when verts changed -> recompute bounding box


def _parse_models(d, o, n, hsize):
    hdr0 = o
    o += n * hsize
    out = []
    for i in range(n):
        b = hdr0 + i * hsize
        m = Model()
        m.hdr = d[b:b + hsize]
        m._nv0 = struct.unpack_from('<H', d, b + 0x6c)[0]
        nv, nf = struct.unpack_from('<HH', d, b + 0x6c)
        nc = struct.unpack_from('<3H', d, b + 0x78)
        if hsize == 0x8c:
            m.x84, nx = struct.unpack_from('<HH', d, b + 0x84)
            m.extra = d[o:o + nx * 2]
            o += nx * 2
        for k in range(nv):
            m.verts.append(struct.unpack_from('<6f', d, o)); o += 24
        for k in range(nf):
            fl, a, b_, c, tex, pad = struct.unpack_from('<6H', d, o)
            m.faces.append(dict(flags=fl, v=(a, b_, c), tex=tex, pad=pad, plane=struct.unpack_from('<4f', d, o + 12)))
            o += 28
        for g in range(3):
            for k in range(nc[g]):
                m.coll[g].append(d[o:o + 32]); o += 32
        out.append(m)
    return out, o


def parse_stpc_models(d):
    """Only the static object models at the start of STPC (u32 n, 1 model each with 0x8c header)."""
    n = struct.unpack_from('<I', d, 0)[0]
    o = 4
    out = []
    for i in range(n):
        ms, o = _parse_models(d, o, 1, 0x8c)
        out += ms
    return out, o


def parse_trak(d):
    n = struct.unpack_from('<I', d, 0)[0]
    models, o = _parse_models(d, 4, n, 0x84)
    assert o == len(d), f'TRAK: {o} != {len(d)}'
    return models


def bbox_header(m):
    if m.verts:
        xs = [v[0] for v in m.verts]; ys = [v[1] for v in m.verts]; zs = [v[2] for v in m.verts]
        lo = (min(xs), min(ys), min(zs)); hi = (max(xs), max(ys), max(zs))
    else:
        lo = hi = (0.0, 0.0, 0.0)
    c = tuple((a + b) / 2 for a, b in zip(lo, hi))
    e = 1 / 128   # the original enlarges the box by 1/128
    h = tuple((b - a) / 2 + e for a, b in zip(lo, hi))
    pts = []
    for s in CORNERS:
        pts += [c[k] + s[k] * h[k] for k in range(3)]
    return c, pts


def build_models(models, hsize, recompute_bbox=None):
    """recompute_bbox=None: keep the original box if present and verts unchanged."""
    hdr, data = b'', b''
    for m in models:
        if recompute_bbox is False or (recompute_bbox is None and m.hdr is not None
                                       and getattr(m, '_nv0', None) == len(m.verts) and not getattr(m, 'dirty', False)):
            h = m.hdr[:0x6c]
        else:
            c, pts = bbox_header(m)
            h = struct.pack('<3f24f', *c, *pts)
        h += struct.pack('<HHII3HHI', len(m.verts), len(m.faces), 0, 0, *(len(g) for g in m.coll), 0, 0)
        if hsize == 0x8c:
            h += struct.pack('<HHI', m.x84, len(m.extra) // 2, 0)
            data += m.extra
        assert len(h) == hsize
        hdr += h
        for v in m.verts:
            data += struct.pack('<6f', *v)
        for f in m.faces:
            data += struct.pack('<6H4f', f['flags'], *f['v'], f['tex'], f.get('pad', 0), *f['plane'])
        for g in m.coll:
            for r in g:
                data += r
    return hdr + data


def build_trak(models):
    return struct.pack('<I', len(models)) + build_models(models, 0x84)


def face_plane(p0, p1, p2):
    ax, ay, az = (p1[i] - p0[i] for i in range(3))
    bx, by, bz = (p2[i] - p0[i] for i in range(3))
    # winding: original normal = (v1-v0) x (v2-v0)
    nx, ny, nz = ay * bz - az * by, az * bx - ax * bz, ax * by - ay * bx
    ln = math.sqrt(nx * nx + ny * ny + nz * nz) or 1.0
    nx, ny, nz = nx / ln, ny / ln, nz / ln
    return (nx, ny, nz, nx * p0[0] + ny * p0[1] + nz * p0[2])


# ---------------------------------------------------------------- collision

def coll_decode(r):
    """32-byte collision polygon -> dict (planes in vertex/export space, i.e. z mirrored back)."""
    typ, mat = r[0], r[1]
    planes = []
    for k in range(5):
        nx, ny, nz, pad, d = struct.unpack_from('<bbbBh', r, 2 + 6 * k)
        planes.append((nx / 127, ny / 127, -nz / 127, d / 4096))
    nedge = 4 if typ & 1 else 3
    return dict(typ=typ, mat=mat, plane=planes[0], edges=planes[1:1 + nedge])


def coll_polygon(c):
    """Corner points from surface plane + edge planes (intersection of adjacent edges)."""
    import numpy as np
    P = c['plane']; E = c['edges']
    pts = []
    n = len(E)
    for i in range(n):
        a, b = E[i], E[(i + 1) % n]
        A = np.array([P[:3], a[:3], b[:3]]); B = np.array([P[3], a[3], b[3]])
        try:
            pts.append(tuple(np.linalg.solve(A, B)))
        except np.linalg.LinAlgError:
            return []
    return pts


def coll_encode(typ, mat, plane, edges):
    """Planes in vertex space -> 32 bytes (z is mirrored for game space)."""
    def q(p):
        nx, ny, nz, d = p
        cl = lambda v: max(-127, min(127, int(round(v * 127))))
        return struct.pack('<bbbBh', cl(nx), cl(ny), cl(-nz), 0, max(-32768, min(32767, int(round(d * 4096)))))
    out = bytes([typ, mat]) + q(plane) + b''.join(q(e) for e in edges)
    return out.ljust(32, b'\0')


def coll_from_polygon(pts, mat=0, typ=None):
    """Convex polygon (3 or 4 points, vertex space, same order as faces) -> 32-byte record."""
    import numpy as np
    P = [np.array(p, float) for p in pts]
    n = np.cross(P[1] - P[0], P[2] - P[0])
    if np.linalg.norm(n) < 1e-9:
        return None
    n /= np.linalg.norm(n)
    d = float(n @ P[0])
    edges = []
    cen = sum(P) / len(P)
    for i in range(len(P)):
        a, b = P[i], P[(i + 1) % len(P)]
        e = np.cross(n, b - a)
        if np.linalg.norm(e) < 1e-9:
            return None
        e /= np.linalg.norm(e)
        if e @ (cen - a) < 0:
            e = -e     # inwards
        edges.append((*e, float(e @ a)))
    if typ is None:
        typ = 1 if len(P) == 4 else 0
    return coll_encode(typ, mat, (*n, d), edges)


# ---------------------------------------------------------------- MAP (tile parts only)
STRAT_REC = 90      # bytes per entry of the 2nd block (0x5c in memory)
NODE_REC = 34       # bytes per entry of the 3rd block (0x30 in memory)


def parse_map_tiles(d, wfpc, vcounts):
    """Returns tile info + byte offsets of the sections (for splicing)."""
    N, W, H = struct.unpack_from('<3I', d, 0)
    o = 12
    offs = {'A': o}
    A = [struct.unpack_from('<3f3I', d, o + 24 * i) for i in range(N)]
    o += 24 * N
    offs['A_end'] = o
    nl = struct.unpack_from('<I', d, o)[0]; o += 4 + 4 * nl
    ns = struct.unpack_from('<I', d, o)[0]; o += 4 + ns * STRAT_REC
    nn = struct.unpack_from('<I', d, o)[0]; o += 4 + nn * NODE_REC
    o += 256
    offs['grid'] = o
    zone = struct.unpack_from('<%dI' % (W * H), d, o); o += 4 * W * H
    offs['B'] = o
    B = [struct.unpack_from('<6i', d, o + 24 * i) for i in range(N)]; o += 24 * N
    midx = list(struct.unpack_from('<%dI' % N, d, o)); o += 4 * N
    anim = []
    if wfpc & 0x10000:
        na = struct.unpack_from('<I', d, o)[0]; o += 4
        anim = [struct.unpack_from('<5I', d, o + 20 * i) for i in range(na)]; o += 20 * na
    offs['col'] = o
    flagged = [i for i in range(N) if A[i][5]]
    altmap = dict(zip(flagged, anim))
    cols, altcols = [], {}
    for i in range(N):
        for k, mdl in enumerate([midx[i]] + ([altmap[i][4]] if i in altmap else [])):
            nv = vcounts[mdl]
            c = list(struct.unpack_from('<%dI' % nv, d, o)); o += 4 * nv
            fr = max([x >> 24 for x in c] + [0])
            frames = [c]
            for f in range(1, fr):
                frames.append(list(struct.unpack_from('<%dI' % nv, d, o))); o += 4 * nv
            if k == 0:
                cols.append(frames)
            else:
                altcols[i] = frames
    offs['col_end'] = o
    return dict(N=N, W=W, H=H, A=A, B=B, midx=midx, anim=anim, zone=zone, cols=cols,
                altcols=altcols, offs=offs)


def build_map_tiles(d, wfpc, old_vcounts, tiles, zones='flat', anim=()):
    """Replaces the tile sections in the MAP chunk.

    tiles: list of dict(x,y,z,rot,model,cols=[[rgba u32]*nv]) ; x/y/z in tile units
    (export space: y up, z<=0), rot in 1/4096 turns (rotation about y).
    Tile swap animations: tiles[i]['flag']=1 + tiles[i]['altcols'] (colours of the alt model);
    anim = list of the 5-u32 entries (cell x, y, z, model, alt model), in the order of the
    flagged tiles. Without anim all flags are cleared.
    zones='flat': disable visibility zones (zone grid = 0, mask of zone 0 = all),
    otherwise ('keep') leave zone grid + 32-entry table unchanged.
    Everything else (strat/node/light blocks) stays unchanged.
    """
    m = parse_map_tiles(d, wfpc, old_vcounts)
    of = m['offs']
    N = len(tiles)
    head = struct.pack('<3I', N, m['W'], m['H'])
    A = b''.join(struct.pack('<3f3I', t['x'], t['y'], t['z'], t['rot'] & 0xfff, t['model'],
                                1 if (anim and t.get('flag')) else 0) for t in tiles)
    B = b''.join(struct.pack('<6i', 0, t['rot'] & 0xfff, 0, int(round(t['x'] * 4096)),
                             int(round(t['y'] * 4096)), int(round(-t['z'] * 4096))) for t in tiles)
    midx = struct.pack('<%dI' % N, *[t['model'] for t in tiles])
    if anim and sum(1 for t in tiles if t.get('flag')) != len(anim):
        raise ValueError('number of flagged tiles != number of anim entries')
    animb = (struct.pack('<I', len(anim)) + b''.join(struct.pack('<5I', *e) for e in anim)) if wfpc & 0x10000 else b''
    cols = b''
    for t in tiles:
        for frame in t['cols'] + (t['altcols'] if (anim and t.get('flag')) else []):
            cols += struct.pack('<%dI' % len(frame), *frame)
    mid = d[of['A_end']:of['B']]
    if zones == 'flat':
        g = of['grid'] - of['A_end']        # offset of zone table+grid within mid
        tab = struct.pack('<2I', 0xffffffff, struct.unpack_from('<I', mid, g - 256 + 4)[0]) + mid[g - 248:g]
        mid = mid[:g - 256] + tab + bytes(4 * m['W'] * m['H'])
    return head + A + mid + B + midx + animb + cols + d[of['col_end']:]


# ---------------------------------------------------------------- TEXT (texture table + pages)

def parse_text(d):
    o = 0
    npal = struct.unpack_from('<I', d, o)[0]; o += 4
    for _ in range(npal):
        n = struct.unpack_from('<I', d, o + 4)[0]; o += 8 + 3 * n
    npg = struct.unpack_from('<I', d, o)[0]; o += 4
    pages = []
    for _ in range(npg):
        fl, w, h = struct.unpack_from('<3I', d, o); o += 12
        if fl == 0:
            sz = w * h; raw = d[o:o + sz]; comp = False
        elif fl & 0x80:
            sz = struct.unpack_from('<I', d, o)[0]; o += 4; raw = d[o:o + sz]; comp = True
        else:
            sz = w * h * 2; raw = d[o:o + sz]; comp = False
        o += sz
        pages.append(dict(flags=fl, w=w, h=h, raw=raw, comp=comp))
    nt = struct.unpack_from('<I', d, o)[0]; o += 4
    texs = []
    for i in range(nt):
        fl, pg, p2, u0, u1, v0, v1 = struct.unpack_from('<HbbBBBB', d, o); o += 8
        texs.append(dict(flags=fl, page=pg, page2=p2, u0=u0, u1=u1, v0=v0, v1=v1))
    return pages, texs


def page_rgba(p):
    """16-bit page -> numpy RGBA (assumption: ABGR1555 like PS1, RLE16 if flags&0x80)."""
    import numpy as np
    w, h = p['w'], p['h']
    if p['flags'] == 0:
        px = np.frombuffer(p['raw'], np.uint8).astype(np.uint16) * 0x421  # greyscale stopgap
    elif p['comp']:
        v = np.frombuffer(p['raw'], '<u2'); out = np.zeros(w * h, np.uint16)
        i = q = 0
        while q < w * h and i < len(v):
            n = int(v[i]); i += 1
            if n & 0x8000:
                n -= 0x10000
            if n > 0:
                out[q:q + n] = v[i:i + n]; i += n; q += n
            elif n < 0:
                out[q:q - n] = v[i]; i += 1; q -= n
        px = out
    else:
        px = np.frombuffer(p['raw'], '<u2')
    px = px[:w * h].reshape(h, w)
    r = (px & 31) << 3; g = ((px >> 5) & 31) << 3; b = ((px >> 10) & 31) << 3
    a = np.where(px == 0, 0, 255)
    return np.dstack([r, g, b, a]).astype(np.uint8)


def tex_uv(t, pw=256, ph=256):
    """UV rectangle (U0,U1,V0,V1) like FUN_00406880 (u1/v1 inclusive -> +1)."""
    return (t['u0'] / pw, (t['u1'] + 1) / pw, t['v0'] / ph, (t['v1'] + 1) / ph)


def face_uvs(flags, rect):
    """Per face corner (v0,v1,v2) the UV from the texture rectangle, per FUN_00556510.

    Bit 0x20 mirrors U; bit 0x800 + bit 0x10 select the triangle half of the quad."""
    U0, U1, V0, V1 = rect
    m = bool(flags & 0x20)
    a, b = (U0, U1) if m else (U1, U0)     # a = "right" without mirroring
    if not (flags & 0x800):
        return [(a, V1), (b, V1), (a, V0)]
    if not (flags & 0x10):
        return [(a, V1), (b, V1), (b, V0)]
    return [(b, V0), (a, V0), (a, V1)]


# ---------------------------------------------------------------- assemble the world

def tile_xform(x, y, z, rot):
    a = rot / 4096 * 2 * math.pi
    ca, sa = math.cos(a), math.sin(a)

    def f(vx, vy, vz):
        return (x + vx * ca - vz * sa, y + vy, z + vx * sa + vz * ca)

    def fn(nx, ny, nz):
        return (nx * ca - nz * sa, ny, nx * sa + nz * ca)
    return f, fn


def load_level(path):
    ch = wad_read(path)
    wf = wfpc_flags(ch)
    trak = parse_trak(wad_get(ch, 'TRAK'))
    vcounts = [len(m.verts) for m in trak]
    mp = parse_map_tiles(wad_get(ch, 'MAP'), wf, vcounts)
    return ch, wf, trak, mp