#!/usr/bin/env python3
"""Parse/write the MAP chunk of level WADs (Lama/Croc 2 engine, PC).

Reimplemented from groove.exe FUN_0042ac50 (MAP loader). Depends on
  WFPC flags (u32 in the WFPC chunk, bits 0x10000, 0x10000000, 0x200000, 0x4000000, 0x10)
  TRAK model vertex counts (vertex colour block).
See notes/objects.md.
"""
import struct


class R:
    def __init__(s, d, o=0):
        s.d, s.o = d, o

    def u32(s):
        v = struct.unpack_from('<I', s.d, s.o)[0]; s.o += 4; return v

    def i32(s):
        v = struct.unpack_from('<i', s.d, s.o)[0]; s.o += 4; return v

    def u16(s):
        v = struct.unpack_from('<H', s.d, s.o)[0]; s.o += 2; return v

    def f(s, fmt):
        v = struct.unpack_from('<' + fmt, s.d, s.o); s.o += struct.calcsize('<' + fmt); return v


def trak_vertcounts(trak):
    """TRAK: u32 n, n x 0x84-Header (u16 @0x6c = vertex count), then data."""
    n = struct.unpack_from('<I', trak, 0)[0]
    return [struct.unpack_from('<H', trak, 4 + i * 0x84 + 0x6c)[0] for i in range(n)]


# field names
TILE = ('x', 'y', 'z', 'rot', 'b', 'alt')                 # 3 f32 + 3 u32
LIGHT = '11I3IH4I2IHHI'                                   # 90 bytes
PATHNODE = 'IIHHH5I'                                      # 34 bytes
STRAT = '3H3iIIi4Iiih' + 'H'                              # 58 bytes (see below)


def parse(d, wfpc, vcounts):
    r = R(d)
    m = {}
    nt, W, H = r.f('3I')
    m['W'], m['H'] = W, H
    m['tiles'] = [r.f('3f3I') for _ in range(nt)]
    n = r.u32(); m['list1'] = list(r.f('%dI' % n))           # referenced by strats via index
    n = r.u32(); m['lights'] = [r.f(LIGHT) for _ in range(n)]
    n = r.u32(); m['nodes'] = [r.f(PATHNODE) for _ in range(n)]
    m['tab32'] = r.f('64I')                                  # 32 x (u32,u32) -> DAT_006da350
    m['grid'] = r.f('%dI' % (W * H))
    m['tiles2'] = [r.f('6I') for _ in range(nt)]
    m['tilemodel'] = list(r.f('%dI' % nt))
    alts = []
    if wfpc & 0x10000:
        n = r.u32(); alts = [r.f('5I') for _ in range(n)]
    m['alts'] = alts
    # vertex colours: nv u32 per tile (byte 3 = frames), more frames if any; then alt model
    altidx = [i for i, t in enumerate(m['tiles']) if t[5]]
    altmap = dict(zip(altidx, alts))
    start = r.o
    split = []                      # vertex colour bytes per tile (incl. alt model)
    for i in range(nt):
        t0 = r.o
        for mdl in [m['tilemodel'][i]] + ([altmap[i][4]] if i in altmap else []):
            nv = vcounts[mdl]
            cols = r.f('%dI' % nv)
            fr = max([c >> 24 for c in cols] + [0])
            if fr > 1:
                r.o += 4 * nv * (fr - 1)
        split.append(d[t0:r.o])
    m['vcol'] = d[start:r.o]
    m['vcol_split'] = split
    ns, x = r.f('2I')
    m['strat_x'] = x
    m['strats'] = [r.f(STRAT) for _ in range(ns)]
    paths = []
    if wfpc & 0x200000:
        np_ = r.u32()
        for _ in range(np_):
            nsub = r.u32(); subs = []
            for _ in range(nsub):
                cnt, fl = r.f('2H'); pts = []
                for _ in range(cnt):
                    a = r.f('2I')
                    c = r.f('I') if wfpc & 0x4000000 else ()
                    pts.append(a + c + r.f('2H'))
                subs.append((fl, pts))
            paths.append(subs)
    m['paths'] = paths
    m['x200'] = r.u32() if wfpc & 0x10 else None
    m['last'] = r.u16()
    m['rest'] = d[r.o:]
    return m


def build(m, wfpc):
    """Counterpart of parse(); m['vcol'] must match tilemodel/alts (raw bytes)."""
    P = struct.pack
    o = [P('<3I', len(m['tiles']), m['W'], m['H'])]
    o += [P('<3f3I', *t) for t in m['tiles']]
    o.append(P('<I', len(m['list1'])) + P('<%dI' % len(m['list1']), *m['list1']))
    o.append(P('<I', len(m['lights'])) + b''.join(P('<' + LIGHT, *x) for x in m['lights']))
    o.append(P('<I', len(m['nodes'])) + b''.join(P('<' + PATHNODE, *x) for x in m['nodes']))
    o.append(P('<64I', *m['tab32']))
    o.append(P('<%dI' % len(m['grid']), *m['grid']))
    o += [P('<6I', *t) for t in m['tiles2']]
    o.append(P('<%dI' % len(m['tilemodel']), *m['tilemodel']))
    if wfpc & 0x10000:
        o.append(P('<I', len(m['alts'])) + b''.join(P('<5I', *a) for a in m['alts']))
    o.append(m['vcol'])
    o.append(P('<2I', len(m['strats']), m['strat_x']))
    o += [P('<' + STRAT, *s) for s in m['strats']]
    if wfpc & 0x200000:
        o.append(P('<I', len(m['paths'])))
        for subs in m['paths']:
            o.append(P('<I', len(subs)))
            for fl, pts in subs:
                o.append(P('<2H', len(pts), fl))
                for p in pts:
                    o.append(P('<2I', *p[:2]) + (P('<I', p[2]) if wfpc & 0x4000000 else b'') + P('<2H', *p[-2:]))
    if wfpc & 0x10:
        o.append(P('<I', m['x200']))
    o.append(P('<H', m['last']))
    o.append(m['rest'])
    return b''.join(o)