#!/usr/bin/env python3
"""Label the spawned-entity records, and MEASURE their stride.

Fixes the two defects recorded in mission-wave-arrivals.md:

  * the previous probe assumed the id chain object+0x04 -> name+0x10 -> char*
    and got '?' for all 116 records. This searches for the chain per record
    instead of assuming one, the way unit_discover.py does.
  * the previous probe assumed RECLEN = 0x200. This measures the gap
    distribution between consecutive records and reports it.

Also samples REMAINING OB so a run can say whether the pilot killed anything --
which is what separates the clock-driven and event-gated wave models.
"""
import os, sys, time, struct, collections, importlib.util

SD = __file__.rsplit('/', 1)[0]
spec = importlib.util.spec_from_file_location('gmem', SD + '/gmem.py')
gmem = importlib.util.module_from_spec(spec); spec.loader.exec_module(gmem)

VT = struct.pack('>I', 0x820AF030)
OB_VA = 0xBDB59668          # REMAINING OB, structures/mission-objective-counter.md
PRINTABLE = set(range(0x20, 0x7F))

def find_records(f, fd, size):
    offs = []
    for start, end in gmem.extents(fd, size):
        pos = start
        while pos < end:
            f.seek(pos); buf = f.read(min(1 << 24, end - pos))
            if not buf: break
            i = buf.find(VT)
            while i != -1:
                if (pos + i) % 4 == 0: offs.append(pos + i)
                i = buf.find(VT, i + 1)
            pos += len(buf)
    return sorted(offs)

def rd(f, off, n):
    f.seek(off); return f.read(n)

def cstr(b):
    e = b.find(b'\x00')
    s = b[:e if e >= 0 else len(b)]
    return s.decode('latin-1') if s and all(c in PRINTABLE for c in s) else None

def resolve_id(f, size, base, scan_words=24, deltas=range(0, 0x41, 4)):
    """Search, don't assume: any pointer in the record's first words, chased
    through one optional indirection, that lands on a UN_/NP_-looking name."""
    head = rd(f, base, scan_words * 4)
    for w in range(0, len(head) - 3, 4):
        (p,) = struct.unpack_from('>I', head, w)
        if not (0x80000000 <= p < 0xC0000000): continue
        o1 = gmem.va_to_off(p)
        if o1 is None or o1 + 0x80 > size: continue
        blk = rd(f, o1, 0x80)
        s = cstr(blk)
        if s and (s.startswith('UN_') or s.startswith('NP_')):
            return s, ('direct', w, 0)
        for d in deltas:                                   # one indirection
            if d + 4 > len(blk): break
            (q,) = struct.unpack_from('>I', blk, d)
            if not (0x80000000 <= q < 0xC0000000): continue
            o2 = gmem.va_to_off(q)
            if o2 is None or o2 + 0x40 > size: continue
            s = cstr(rd(f, o2, 0x40))
            if s and (s.startswith('UN_') or s.startswith('NP_')):
                return s, ('indirect', w, d)
    return None, None

def ob(f):
    o = gmem.va_to_off(OB_VA)
    if o is None: return None
    return struct.unpack('>I', rd(f, o, 4))[0]

def main():
    secs = int(sys.argv[1]) if len(sys.argv) > 1 else 120
    every = int(sys.argv[2]) if len(sys.argv) > 2 else 10
    path = gmem.mem_path()
    fd = os.open(path, os.O_RDONLY); size = os.fstat(fd).st_size
    f = os.fdopen(os.dup(fd), 'rb')
    offs = find_records(f, fd, size)
    if not offs:
        print('NO ENTITY RECORDS -- not in a mission'); return 2
    print('records: %d' % len(offs))

    gaps = [offs[i+1] - offs[i] for i in range(len(offs) - 1)]
    hist = collections.Counter(gaps)
    print('\n--- MEASURED stride (gap between consecutive records) ---')
    print('  min=%d  median=%d  most common: %s'
          % (min(gaps), sorted(gaps)[len(gaps)//2],
             [(g, c) for g, c in hist.most_common(6)]))

    print('\n--- label resolution (searched, not assumed) ---')
    labels, how = {}, collections.Counter()
    for k, o in enumerate(offs):
        s, h = resolve_id(f, size, o)
        labels[k] = s
        how[h[0] if h else 'FAILED'] += 1
        if h: how[('chain', h[1], h[2])] += 1
    named_n = sum(1 for v in labels.values() if v)
    print('  resolved %d/%d' % (named_n, len(offs)))
    print('  by method:', [(str(a), b) for a, b in how.most_common(6)])
    if named_n:
        print('  unit id histogram:',
              collections.Counter(v for v in labels.values() if v).most_common(10))
    else:
        print('  STILL UNRESOLVED -- the id is not reachable from the record head')

    print('\n--- REMAINING OB over %ds (does the pilot kill anything?) ---' % secs)
    t0 = time.time(); seen = []
    while time.time() - t0 < secs:
        v = ob(f); seen.append((round(time.time() - t0), v))
        print('  t=%4ds  OB=%s' % seen[-1], flush=True)
        time.sleep(every)
    vals = [v for _, v in seen if v is not None]
    print('  distinct OB values: %s' % sorted(set(vals))[:10])
    return 0

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