#!/usr/bin/env python3
"""REMAINING OB by HUD value, on the shared harness.

Same idea as ob_by_hud.py -- read the counter off screen, keep heap positions
equal to it, intersect across readings -- but the witness, the incremental save
and the baseline discard now come from probeharness.Probe rather than being
re-implemented (and, twice, forgotten).

Candidates persist to /tmp/ob_candidates.json after EVERY reading, so a chained
attach continues the same intersection on the same mission. Offsets are only
valid within one emulator instance; the session clears the file at launch.
"""
import json, os, subprocess, sys
sys.path.insert(0, __file__.rsplit('/', 1)[0])
from probeharness import Probe
import gmem
import numpy as np
import ob_read

SHOT = '/tmp/ob_hud.png'
CAND = '/tmp/ob_candidates.json'
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 hud():
    subprocess.run(['screenshot', SHOT], capture_output=True, timeout=60)
    try:
        txt, _ = ob_read.read(SHOT)
    except Exception as e:
        return None, str(e)
    t = (txt or '').strip()
    return (int(t), 'ok') if t.isdigit() else (None, 'unreadable %r' % txt)

def matches(buf, value):
    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)
        out[name] = (np.nonzero(a == value)[0].astype(np.int64) * w + skew)
    return out

def main():
    secs  = int(sys.argv[1]) if len(sys.argv) > 1 else 300
    every = int(sys.argv[2]) if len(sys.argv) > 2 else 25
    p = Probe(baseline=116)
    if not p.ok:
        print(p.why); return 3
    print(p.summary(), flush=True)
    p.log('t\thud\tencoding\tcandidates\tstalled', '/tmp/ob_hunt2.tsv')

    cand, base, seen = None, None, []
    if os.environ.get('SYLPH_OB_RESUME') == '1' and os.path.exists(CAND):
        raw = json.load(open(CAND))
        cand = {k: np.array(v, dtype=np.int64) for k, v in raw['cand'].items()}
        seen = raw['seen']
        print('resumed: %s' % {k: len(v) for k, v in cand.items()}, flush=True)

    while True:
        v, note = hud()
        buf, base = p.heap()
        st = p.status()
        if v is None:
            print('t=%4ds  HUD unreadable (%s)%s' % (p.elapsed, note, st), flush=True)
        else:
            hit = matches(buf, v)
            cand = hit if cand is None else {
                k: np.intersect1d(cand[k], hit[k], assume_unique=True)
                for k in cand if k in hit}
            seen.append(v)
            json.dump({'cand': {k: v2.tolist() for k, v2 in cand.items()},
                       'seen': seen}, open(CAND, 'w'))
            desc = '  '.join('%s:%d' % (k, len(cand[k])) for k in cand)
            print('t=%4ds  HUD=%-4d  %s%s' % (p.elapsed, v, desc, st), flush=True)
            for k in cand:
                p.emit('%d\t%d\t%s\t%d\t%s' % (p.elapsed, v, k, len(cand[k]), st.strip()))
        if not p.tick(every, secs):
            break

    print('\n%s' % p.summary())
    print('HUD values seen: %s' % sorted(set(seen)))
    if cand:
        for k in sorted(cand, key=lambda k: len(cand[k])):
            c = cand[k]
            print('  %-8s %d' % (k, len(c)))
            if 0 < len(c) <= 8 and len(set(seen)) >= 2:
                for off in c:
                    va = gmem.primary_va(base + int(off))
                    print('      va %s' % (('%#010x' % va) if va else '?'))
    return 0

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