#!/usr/bin/env python3
"""Strat types (object types) of a level WAD: names from the STPC strat code, MAP instances.

Strat code sits at the end of STPC (after models/animsets). Bytecode = u32 words
(low16 = opcode, high16 = short operand). A strat (almost always) starts with
  ... 0xd4, 0xea, <s32 rel>  -> name = C string at (address after rel) + rel
  strats.py level.wad [--all]
"""
import os
import struct
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import map_dump  # noqa: E402
import stpcwalk  # noqa: E402


def cstr(d, o):
    e = d.index(b'\0', o)
    return d[o:e].decode('latin1')


def strat_name(st, entry, lo):
    """Name of a MAP entry: spawn stub (b2 <abs> e0 ...) -> target; otherwise first d4/ea from entry."""
    w = struct.unpack_from('<8I', st, entry)
    for k in range(5):
        if w[k] == 0xb2 and w[k + 2] == 0xe0 and lo <= w[k + 1] < len(st):
            n = strat_name(st, w[k + 1], lo)
            return n + '(spawn)' if n else None
    for k in range(0, 1024, 4):
        o = entry + k
        if o + 12 > len(st):
            break
        a, b = struct.unpack_from('<II', st, o)
        if a == 0xd4 and b == 0xea:
            rel = struct.unpack_from('<i', st, o + 8)[0]
            n = o + 12 + rel
            if lo <= n < len(st):
                return cstr(st, n)
    return None


def all_named(st, lo):
    """All d4/ea name sites in the code region (word grid from lo)."""
    out = []
    for o in range(lo, len(st) - 12, 4):
        a, b = struct.unpack_from('<II', st, o)
        if a == 0xd4 and b == 0xea:
            rel = struct.unpack_from('<i', st, o + 8)[0]
            n = o + 12 + rel
            if lo <= n < o:
                out.append((o, cstr(st, n)))
    return out


def level_types(path):
    ch, wf, m = map_dump.load(path)
    st = ch['STPC']
    lo = stpcwalk.walk(st, wf)['code']
    types = {}
    for s in m['strats']:
        code = s[6]
        nm = strat_name(st, code, lo) or f'?@{code:#x}'
        types.setdefault(nm, []).append(s)
    return ch, wf, m, lo, types


if __name__ == '__main__':
    ch, wf, m, lo, types = level_types(sys.argv[1])
    print(f'code region {lo:#x}..{len(ch["STPC"]):#x}')
    for nm, ss in sorted(types.items()):
        s = ss[0]
        print(f'{len(ss):3d} x {nm:32s} code@{s[6]:#x} nparams={s[7]} nvars={s[9]} f30={s[10]} '
              f'flags={s[11]:#x} extra={s[12]:#x}')
    if '--all' in sys.argv:
        used = {s[6] for v in types.values() for s in v}
        print('-- all named strats in the code (name op address):')
        for o, n in all_named(ch['STPC'], lo):
            print(f'  {o:#x} {n}')