#!/usr/bin/env python3
"""Decode a stage's squadron roster: hidden/dat `stage\\UnitGroup_S<NN>.tbl`.

Pure static work — reads the extracted disc, runs no emulator.

    python3 tools/re-capture/unitgroup.py S02
    python3 tools/re-capture/unitgroup.py --all --check

Format (see docs/re/structures/unit-group-table.md):

    "IDXD"                      magic
    u32 nrec                    record count; exactly one record is the
                                Enumerate_Squadrons roster, not a squadron
    nrec x 16                   (key, squadron_string_off, field_lo, field_hi)
    u32 npool                   property-entry count
    npool x 12                  (tag, name_off | 0xffffffff, value_off)
    u32 strsize                 string-pool byte size; STR+strsize == len(file)
    strsize bytes               string pool; [0] is "Enumerate_Squadrons",
                                [0x14] is the empty string. Every offset in
                                the file -- squadron ids included -- is relative
                                to STR.

A squadron's field list is `Count` member tuples of four positional entries
(unit model, message set, a number, pilot character) followed by five named
entries: Count, SideID, AIID, FormationID, DisableInterval. So

    len(fields) == Count * 4 + 5

which holds for 1160 of 1160 squadrons across the 28 stage tables present in
GP_MAIN_GAME_E.pak.
"""
import argparse, glob, os, struct, sys, zlib

MODULUS, RECIP = 0x00FFF9D7, 0x80031493

def _rotl(v, n): return ((v << n) | (v >> (32 - n))) & 0xFFFFFFFF

def name_hash(s):
    bs = bytearray(s.encode('latin-1', 'replace'))
    for i, x in enumerate(bs):
        if 65 <= x <= 90: bs[i] = x + 0x20
    a = b = 0
    for byte in bs:
        c = ((byte - 256) if byte > 127 else byte) & 0xFFFFFFFF
        a = ((_rotl(a, 8) & 0xFFFFFF00) + c) & 0xFFFFFFFF
        b = (b + c) & 0xFFFFFFFF
        q = _rotl(((a * RECIP) >> 32) & 0xFFFFFFFF, 9) & 0x1FF
        a = (a - (q * MODULUS)) & 0xFFFFFFFF
    return (((b << 24) & 0xFF000000) | (a & 0x00FFFFFF)) & 0xFFFFFFFF

TAG_MODULUS = (1 << 24) - 33      # 0x00FFFFDF, prime

def tag_hash(s):
    """IDXD record key / field tag -- NOT name_hash.

    Same shape as name_hash (8-bit byte-sum checksum over a 24-bit modular
    polynomial) but modulo 0x00FFFFDF instead of 0x00FFF9D7, and NOT
    lowercased, so tags are case-sensitive. Recovered empirically from the 8643
    (name -> tag) pairs the tables themselves carry; `unitgroup.py --checktags`
    re-verifies all of them. A record's key is the tag of its own name, which
    each table lists in an in-table roster record.
    """
    b = s.encode()
    lo = 0
    for c in b:
        lo = (lo * 256 + c) % TAG_MODULUS
    return ((sum(b) & 0xFF) << 24) | lo

def read_entry(pak, h):
    idx = open(pak, 'rb').read()
    base = pak[:-4]
    data = b''.join(open(p, 'rb').read() for p in sorted(glob.glob(base + '.p[0-9][0-9]')))
    for i in range(struct.unpack_from('>I', idx, 4)[0]):
        k, off, csize = struct.unpack_from('>III', idx, 16 + i * 12)
        if k == h:
            s = data[off:off + csize]
            return zlib.decompress(s[10:]) if s[:2] == b'Z1' else s
    raise KeyError('%#010x not in %s' % (h, os.path.basename(pak)))

def parse(b):
    u = lambda o: struct.unpack_from('>I', b, o)[0]
    if b[:4] != b'IDXD': raise ValueError('bad magic')
    nrec = u(0x04)
    npool_off = 0x08 + nrec * 16
    npool = u(npool_off)
    pool = npool_off + 4
    strsize_off = pool + npool * 12
    STR = strsize_off + 4
    if STR + u(strsize_off) != len(b):
        raise ValueError('string-pool size trailer mismatch')
    def s(o):
        return b[o:b.index(b'\x00', o)].decode('latin-1')
    out = []
    for i in range(nrec):
        key, sq, lo, hi = (u(0x08 + i * 16 + j * 4) for j in range(4))
        fields = []
        for j in range(hi - lo):
            o = pool + (lo + j) * 12
            tag, name, val = u(o), u(o + 4), u(o + 8)
            fields.append((tag, s(STR + name) if name != 0xFFFFFFFF else None, s(STR + val)))
        out.append({'key': key, 'squadron': s(STR + sq), 'fields': fields})
    return out

def named(rec):   return {n: v for _, n, v in rec['fields'] if n is not None}
def members(rec):
    pos = [v for _, n, v in rec['fields'] if n is None]
    return [tuple(pos[i:i + 4]) for i in range(0, len(pos), 4)]
def is_roster(rec):  return 'Count' not in named(rec)

def roster(recs):
    """One record enumerates key -> squadron id; use it to cross-check the rest.

    It is usually the last record but not always (S05, S28, ... put it earlier),
    so find it by its lack of a Count field rather than by position.
    """
    for r in recs:
        if is_roster(r):
            return {tag: n for tag, n, _ in r['fields']}
    return {}

def load_stage(pak, stage):
    return parse(read_entry(pak, name_hash('stage\\UnitGroup_%s.tbl' % stage)))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('stage', nargs='?', default='S02')
    ap.add_argument('--pak', default='/work/sylph_extract/dat/GP_MAIN_GAME_E.pak')
    ap.add_argument('--all', action='store_true', help='every stage S01..S29')
    ap.add_argument('--check', action='store_true', help='only report the Count*4+5 identity')
    a = ap.parse_args()
    stages = ['S%02d' % n for n in range(1, 30)] if a.all else [a.stage]
    held = total = 0
    for st in stages:
        try:
            recs = load_stage(a.pak, st)
        except KeyError as e:
            print('%s: absent (%s)' % (st, e)); continue
        squads = [r for r in recs if not is_roster(r)]
        ok = sum(1 for r in squads if len(r['fields']) == int(named(r)['Count']) * 4 + 5)
        rst = roster(recs)
        agree = sum(1 for r in squads if rst.get(r['key']) == r['squadron'])
        held += ok; total += len(squads)
        print('%s: %3d squadrons, %4d members, Count*4+5 %d/%d, roster agrees %d/%d'
              % (st, len(squads), sum(len(members(r)) for r in squads), ok, len(squads),
                 agree, len(squads)))
        if a.check: continue
        for r in squads:
            d = named(r)
            print('  %-10s %-8s Count=%-3s DisableInterval=%-4s %-30s %s'
                  % (r['squadron'], d['SideID'], d['Count'], d['DisableInterval'],
                     d['FormationID'], d['AIID']))
            for m in members(r):
                print('      unit=%-34s msg=%-24s n=%-4s pilot=%s' % m)
    if len(stages) > 1:
        print('\nTOTAL %d squadrons, identity held %d/%d' % (total, held, total))
    return 0 if held == total else 1

if __name__ == '__main__':
    sys.exit(main())