#!/usr/bin/env python3
"""Locate REMAINING OB from the HUD value instead of from kill events.

The correlation route works but needs marked-fighter kills, and the pilot gets
about two per five minutes (remaining-ob-hunt.md). The HUD already shows the
number, and ob_read.py already reads it, so match the displayed value against
memory directly -- no kills required.

Each reading intersects: keep heap words equal to the value the HUD shows at
that moment. A second reading at a DIFFERENT value collapses the set hard; even
repeated readings at the same value help, since unrelated words drift.
"""
import os, sys, time, subprocess, collections
sys.path.insert(0, __file__.rsplit('/', 1)[0])
import gmem, gworld
import numpy as np
import ob_read

LO, HI = 0xBD000000, 0xBE000000
SHOT = '/tmp/ob_hud.png'

def region(fd):
    lo, hi = gmem.va_to_off(LO), gmem.va_to_off(HI)
    out, pos = bytearray(), lo
    while pos < hi:
        n = min(1 << 24, hi - pos); out += os.pread(fd, n, pos); pos += n
    return bytes(out), lo

# The big-endian u32 reading was refuted: no word held 4 then 11. Widen rather
# than assume -- the counter may be narrower, little-endian, or unaligned. Each
# encoding keeps its own candidate set, expressed as BYTE OFFSETS so the answer
# is directly usable whichever one wins.
ENCODINGS = [
    ('u32be', '>u4', 4, 0), ('u32le', '<u4', 4, 0),
    ('u16be', '>u2', 2, 0), ('u16be@1', '>u2', 2, 1),
    ('u16le', '<u2', 2, 0), ('u16le@1', '<u2', 2, 1),
    ('u8',    'u1',  1, 0),
]

def matches(buf, value):
    """byte offsets whose value equals `value`, per encoding"""
    out = {}
    for name, dt, w, skew in ENCODINGS:
        if value > (1 << (8 * w)) - 1:
            continue
        n = (len(buf) - skew) // w * w
        a = np.frombuffer(buf[skew:skew + n], dtype=dt)
        idx = np.nonzero(a == value)[0]
        out[name] = idx.astype(np.int64) * w + skew
    return out

def hud_value():
    subprocess.run(['screenshot', SHOT], capture_output=True, timeout=60)
    try:
        txt, scores = ob_read.read(SHOT)
    except Exception as e:
        return None, str(e)
    t = (txt or '').strip()
    if not t.isdigit(): return None, 'unreadable %r' % txt
    return int(t), 'scores %s' % (scores,)   # ob_read returns tuples, not floats

def main():
    secs  = int(sys.argv[1]) if len(sys.argv) > 1 else 240
    every = int(sys.argv[2]) if len(sys.argv) > 2 else 25
    w = gworld.World(); fd = w.fd
    # Witness. This is the FOURTH probe written without one, and the third to
    # produce a run whose flat output could not be distinguished from a freeze.
    # The recurring fix is a shared harness; until that exists, carry it.
    b0, _ = region(fd); time.sleep(3.0); b1, _ = region(fd)
    a0 = np.frombuffer(b0, dtype='>u4').astype(np.int64)
    a1 = np.frombuffer(b1, dtype='>u4').astype(np.int64)
    d = a1 - a0
    rates = collections.Counter((d[(d > 15) & (d < 600)] // 3).tolist())
    band = [r for r in rates if 8 <= r <= 40]
    pick = max(band, key=lambda r: rates[r]) if band else None
    ticks = np.nonzero(d // 3 == pick)[0][:32] if pick is not None else np.array([], dtype=int)
    print('tick witnesses: %d at %s/s' % (len(ticks), pick), flush=True)
    last_t = a1[ticks] if len(ticks) else None
    cand = None; base = None; seen = []
    t0 = time.time()
    while time.time() - t0 < secs:
        v, note = hud_value()
        el = round(time.time() - t0)
        if v is None:
            print('  t=%4ds  HUD unreadable (%s)' % (el, note), flush=True)
        else:
            buf, base = region(fd)
            hit = matches(buf, v)
            if cand is None:
                cand = hit
            else:
                cand = {k: np.intersect1d(cand[k], hit[k], assume_unique=True)
                        for k in cand if k in hit}
            seen.append(v)
            st = ''
            if len(ticks):
                now_t = np.frombuffer(buf, dtype='>u4').astype(np.int64)[ticks]
                if int((now_t > last_t).sum()) == 0: st = '  *** GUEST STALLED ***'
                last_t = now_t
            print('  t=%4ds  HUD=%-4d  %s%s  (%s)'
                  % (el, v, '  '.join('%s:%d' % (k, len(cand[k])) for k in cand), st, note),
                  flush=True)
            live = {k: c for k, c in cand.items() if len(c)}
            if live and len(set(seen)) >= 2 and min(len(c) for c in live.values()) <= 40:
                break
        time.sleep(every)
    print('\nHUD values seen: %s' % sorted(set(seen)))
    if cand and base is not None:
        for k in sorted(cand, key=lambda k: len(cand[k])):
            c = cand[k]
            print('  %-8s %d candidate(s)' % (k, len(c)))
            for off in c[:8]:
                va = gmem.primary_va(base + int(off))
                print('      va %s' % (('%#010x' % va) if va else '?'))
        if all(len(c) == 0 for c in cand.values()):
            print('\nEVERY encoding eliminated -- the counter is not in this region '
                  'in any of them, or a HUD reading is wrong.')
    else:
        print('no candidates (HUD never read)')
    return 0

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