#!/usr/bin/env python3
"""Dump any `stage\\*.tbl` mission table from GP_MAIN_GAME_<lang>.pak.

Pure static work -- reads the extracted disc, runs no emulator. The container is
the same self-describing IDXD format decoded for the squadron roster; see
docs/re/structures/unit-group-table.md and tools/re-capture/unitgroup.py.

    python3 tools/re-capture/stagetbl.py Stage_S02          # by stage record
    python3 tools/re-capture/stagetbl.py --name 'stage\\AIParams_S02.tbl'
    python3 tools/re-capture/stagetbl.py --follow S02       # stage record + all
                                                           # tables it names
"""
import argparse, re, sys
sys.path.insert(0, __file__.rsplit('/', 1)[0])
from unitgroup import read_entry, name_hash, parse, named

PAK = '/work/sylph_extract/dat/GP_MAIN_GAME_E.pak'

def stage_record_hash(stage):
    """The per-stage definition record is not name-addressed; find it by content."""
    import struct, glob, os, zlib
    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]')))
    want = ('UnitGroup_%s.tbl' % stage).encode()   # the _Test template names
                                                   # Stage_S02.xpr too, so key
                                                   # off the squadron table
    for i in range(struct.unpack_from('>I', idx, 4)[0]):
        h, off, cs = struct.unpack_from('>III', idx, 16 + i * 12)
        s = data[off:off + cs]
        try: b = zlib.decompress(s[10:]) if s[:2] == b'Z1' else s
        except Exception: continue
        if b[:4] == b'IDXD' and want in b and b'EnumerateSquadron' in b:
            return h
    raise KeyError('no stage record for ' + stage)

def dump(label, recs, limit=None):
    print('=== %s (%d records) ===' % (label, len(recs)))
    for r in recs[:limit]:
        print('-- %r (%d fields)' % (r['squadron'], len(r['fields'])))
        for tag, fn, v in r['fields']:
            print('    %08x %-30s = %s' % (tag, fn if fn else '(pos)', v))
    if limit and len(recs) > limit:
        print('    ... %d more records not shown' % (len(recs) - limit))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('target', nargs='?', default='Stage_S02')
    ap.add_argument('--name', help=r'archive name, e.g. stage\AIParams_S02.tbl')
    ap.add_argument('--follow', metavar='SNN', help='stage record plus every .tbl it names')
    ap.add_argument('--limit', type=int, default=None)
    a = ap.parse_args()
    if a.name:
        dump(a.name, parse(read_entry(PAK, name_hash(a.name))), a.limit); return 0
    stage = a.follow or (a.target[6:] if a.target.startswith('Stage_') else a.target)
    recs = parse(read_entry(PAK, stage_record_hash(stage)))
    dump('Stage_%s definition record' % stage, recs)
    if not a.follow: return 0
    seen = set()
    for r in recs:
        for _, fn, v in r['fields']:
            if not v.endswith('.tbl') or v in seen: continue
            seen.add(v)
            src, n = PAK, 'stage\\' + v.split('+')[-1]
            if '+' in v and 'DefTables' in v:
                src, n = '/work/sylph_extract/hidden/DefTables.pak', v.split('+')[-1]
            print()
            try: dump('%s  (via %s)' % (n, fn), parse(read_entry(src, name_hash(n))), a.limit)
            except Exception as e: print('=== %s -> unresolved: %s' % (n, e))
    return 0

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