#!/usr/bin/env python3
"""Census the SUB-RECORDS inside the 114 unit tables, and check the counts.

A unit `.tbl` is one IDXD pak entry.  Besides `Generic`/`Maneuver`/`Effect` it
carries a `StructureCount` record and a variable number of destructible-part
records: `Turret_NNN`, `Bridge_NNN`, `Hatch_NNN`, `ShieldGenerator_NNN`,
`Thruster_NNN`, `Versatile_NNN`.  Regenerates docs/re/data/unit-substructures.txt.
"""
import sys, os, glob, re, collections
from disc import disc_root

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import unitgroup as U

# pak_entries lives in the scratch decoder; fall back to a local copy of the recipe
import struct, zlib
def pak_entries(path):
    idx = open(path, 'rb').read()
    n = struct.unpack_from('>I', idx, 4)[0]
    parts = sorted(glob.glob(re.sub(r'\.pak$', '', path) + '.p[0-9][0-9]'))
    data = b''.join(open(p, 'rb').read() for p in parts)
    out = []
    for i in range(n):
        h, off, csz = struct.unpack_from('>III', idx, 16 + i * 12)
        s = data[off:off + csz]
        if s[:2] == b'Z1':
            try: s = zlib.decompress(s[10:])
            except Exception: continue
        out.append((h, s))
    return out

KIND = {'Turret': 'TurretCount', 'Bridge': 'BridgeCount', 'Hatch': 'HatchCount',
        'ShieldGenerator': 'ShieldGeneratorCount', 'Thruster': 'ThrusterCount',
        'Versatile': 'VersatileCount'}

def blank(n):
    return not n.get('Name') and not n.get('NomalModel') and not n.get('Frame')

def units_of(pak):
    out = []
    for h, b in pak_entries(pak):
        if b[:4] != b'IDXD': continue
        try: recs = U.parse(b)
        except Exception: continue
        g = [r for r in recs if r['squadron'] == 'Generic']
        if g and 'HP' in U.named(g[0]):
            out.append((recs, U.named(g[0])))
    return out

def main():
    pak = glob.glob(disc_root() + '/**/GP_MAIN_GAME_E.pak', recursive=True)[0]
    units = units_of(pak)
    print("# Sub-records inside the %d unit tables of %s" % (len(units), os.path.basename(pak)))
    print("# Regenerate: python3 tools/re-capture/unit_substructures.py")
    print("# See docs/re/structures/unit-substructure-records.md")

    norm = lambda n: re.sub(r'_\d{3}$', '_NNN', n)
    fam = collections.Counter(); per = collections.Counter()
    core = collections.defaultdict(collections.Counter)
    for recs, gn in units:
        seen = collections.Counter()
        for r in recs:
            n = norm(r['squadron']); fam[n] += 1; seen[n] += 1
            for _, fn, _ in r['fields']:
                if fn and not re.search(r'_\d{3}$', fn): core[n][fn] += 1
        for n, c in seen.items(): per[n] = max(per[n], c)
    print("\n## families (records, max per unit, distinct non-indexed field names)")
    for n, c in fam.most_common():
        print("   %-22s %5d   max/unit %3d   fields %3d" % (n, c, per[n], len(core[n])))

    for n, _ in fam.most_common():
        if n in ('Generic', 'Maneuver'): continue
        print("\n## %s -- non-indexed fields (records carrying it, of %d)" % (n, fam[n]))
        for k, v in sorted(core[n].items(), key=lambda x: (-x[1], x[0])):
            print("     %-32s %5d" % (k, v))

    res = collections.Counter(); exc = collections.defaultdict(list)
    for recs, gn in units:
        sc = U.named([r for r in recs if r['squadron'] == 'StructureCount'][0])
        for k, f in KIND.items():
            subs = [r for r in recs if re.fullmatch(k + r'_\d{3}', r['squadron'])]
            w = sc.get(f, '').strip()
            if not w.isdigit(): res['no field']+= 1; continue
            d = int(w)
            if d == len(subs): res['equal']+= 1
            elif d == 0 and len(subs) == 1 and blank(U.named(subs[0])):
                res['0 declared, 1 blank placeholder'] += 1
            else:
                res['DIFFERS'] += 1
                exc[(k, d, len(subs))].append(gn.get('ID'))
    print("\n## CONTROL  StructureCount.<Kind>Count == #<Kind>_NNN records")
    print("##          (a 0 count still emits one blank slot)")
    print("   %d comparisons (%d units x %d kinds)" % (len(units) * len(KIND), len(units), len(KIND)))
    for k in ('equal', '0 declared, 1 blank placeholder', 'DIFFERS', 'no field'):
        print("     %-34s %4d" % (k, res[k]))
    print("   exceptions:")
    for kk in sorted(exc, key=lambda x: -len(exc[x])):
        print("     kind=%-8s declared=%-3d records=%-3d x%d" % (kk[0], kk[1], kk[2], len(exc[kk])))
        for i in sorted(exc[kk]): print("        %s" % i)

    wids = collections.Counter(); wnames = set()
    for h, b in pak_entries(pak):
        if b[:4] != b'IDXD': continue
        try: recs = U.parse(b)
        except Exception: continue
        for r in recs:
            if r['squadron'] == 'Weapon':
                i = U.named(r).get('ID')
                if i: wnames.add(i)
    for recs, gn in units:
        for r in recs:
            if re.fullmatch(r'Turret_\d{3}', r['squadron']):
                wids[U.named(r).get('WeaponID', '')] += 1
    hit = sum(v for k, v in wids.items() if k in wnames)
    print("\n## CONTROL  Turret_NNN.WeaponID resolves into the Weapon datasheet")
    print("   Weapon.ID values: %d" % len(wnames))
    print("   turrets: %d, distinct WeaponID: %d" % (sum(wids.values()), len(wids)))
    print("   resolve: %d   do NOT resolve: %d" % (hit, sum(wids.values()) - hit))
    print("   weapons never mounted on a turret: %d" % len(wnames - set(wids)))

if __name__ == "__main__":
    main()