#!/usr/bin/env python3
"""Is `REMAINING OB` the population of a per-entity *bit*?

[`ob_flag.py`](ob_flag.py) refuted the word-level version: no 4-byte word near an
entity's position holds a common value on exactly the objective entities, and its
two candidates both died on the `12 → 11` transition. But that test asks which
entities share an **exact 32-bit value**, and it names its own blind spot — a
single bit ORed into a word that also carries health, a timer or a state machine
would never produce a shared value and could not be seen.

This is that test at bit granularity: for every 4-byte offset in the window and
every one of its 32 bits, count how many entities have the bit SET, keep the
(offset, bit) pairs whose count is exactly the counter, and then require the
survivors to match again after the counter moves — the same "verify across a
transition you did not select on" rule, because there will be thousands of
coincidences at N alone.

Both polarities are counted: an objective might be marked by a bit that is set on
it, or by one that is CLEAR on it and set on everything else.

🔴 Same limits as the word-level test, and they still bound the answer:
`entities2.typed` types entities by their position CHANGING, so a stationary
objective is invisible; and anything outside the window is untested.

Usage: ob_bitflag.py <out-dir> [timeout_s]
"""
import json
import os
import struct
import subprocess
import sys
import time

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import entities2  # noqa: E402
import frozen     # noqa: E402
import gmem       # noqa: E402
import gworld     # noqa: E402
import ob_read    # noqa: E402

# The counter's address is per-run, but not arbitrary: every run measured so far
# has put it in a narrow band, and three exact addresses have now been confirmed
# or strongly implicated. Trying the known ones against the HUD costs nothing and
# no transitions, where ob_hunt.py needs two — and with roughly half of all runs
# ending early, transitions are the scarce resource.
KNOWN_VAS = [0xBDB59668, 0xBDB49668, 0xBDB58668]
VA = int(os.environ["OB_VA"], 0) if "OB_VA" in os.environ else None
LO, HI = -0x400, 0xC00          # window around the position triple
DELTA = 0x130


def counter(fd, va=None):
    return struct.unpack(">I", os.pread(fd, 4, gmem.va_to_off(va or VA)))[0]


def locate(fd, hud_value):
    """Pick whichever known address currently equals the HUD, or None."""
    for va in ([VA] if VA else KNOWN_VAS):
        try:
            if counter(fd, va) == hud_value:
                return va
        except Exception:
            continue
    return None


def hud(shot):
    subprocess.run(["screenshot", shot], stdout=subprocess.DEVNULL,
                   stderr=subprocess.DEVNULL)
    txt, _ = ob_read.read(shot)
    return int(txt) if txt.isdigit() else None


def sample(w, defs, tag):
    """(counter, entity names, words[n_entities, n_words] as uint32)."""
    fd = w.fd
    n = counter(fd)
    movers = entities2.moving(fd, w.size)
    ents = entities2.typed(fd, defs, movers, DELTA)
    uniq = {}
    for off, nm, pos, sp in ents:
        uniq.setdefault((nm, tuple(round(c, 1) for c in pos)), (off, nm))
    ents = list(uniq.values())
    rows, names = [], []
    for off, nm in ents:
        blob = os.pread(fd, HI - LO, off + LO)
        if len(blob) != HI - LO:
            continue
        rows.append(np.frombuffer(blob, dtype=">u4"))
        names.append(nm)
    print(f"[{tag}] counter={n} entities={len(rows)}", flush=True)
    return n, names, (np.array(rows, dtype=np.uint32) if rows else np.zeros((0, 0), np.uint32))


def bit_counts(words):
    """counts[bit, word] = how many entities have that bit set."""
    out = np.zeros((32, words.shape[1]), dtype=np.int32)
    for b in range(32):
        out[b] = ((words >> np.uint32(b)) & np.uint32(1)).sum(axis=0)
    return out


def main():
    out = sys.argv[1]
    deadline = time.time() + (float(sys.argv[2]) if len(sys.argv) > 2 else 700)
    os.makedirs(out, exist_ok=True)
    w = gworld.World()
    defs = entities2.definitions(w)

    global VA
    v = None
    for _ in range(8):
        v = hud(f"{out}/a.png")
        if v is not None:
            break
        time.sleep(3)
    if v is None:
        print("could not read the HUD counter at all — is this in flight?",
              flush=True)
        return 2
    found = locate(w.fd, v)
    if found is None:
        tried = " ".join(f"{a:#x}" for a in ([VA] if VA else KNOWN_VAS))
        print(f"HUD={v} but none of the known addresses holds it ({tried}) — "
              f"re-scan with ob_hunt.py and pass OB_VA", flush=True)
        return 2
    VA = found
    n0 = counter(w.fd)
    print(f"HUD={v} RAM={n0} at {VA:#x}", flush=True)

    for _ in range(6):
        nA, namesA, wordsA = sample(w, defs, "A")
        if len(namesA) >= 10:
            break
        print("[A] too few entities — retrying", flush=True)
        time.sleep(5)
    if len(namesA) < 10:
        print("[A] never got a usable entity sample", flush=True)
        return 4

    cA = bit_counts(wordsA)
    nA_ents = len(namesA)
    setA = np.argwhere(cA == nA)                 # bit set on exactly N entities
    clrA = np.argwhere(cA == nA_ents - nA)       # bit CLEAR on exactly N
    print(f"[A] (offset,bit) pairs set on exactly {nA}: {len(setA)}; "
          f"clear on exactly {nA}: {len(clrA)}", flush=True)

    stuck = 0
    while time.time() < deadline:
        time.sleep(5)
        if counter(w.fd) != nA:
            break
        stuck += 1
        if stuck % 12 == 0:
            # Two different ways the wait can be pointless, and only one of them
            # is a freeze: a run that ends in GAME OVER keeps ANIMATING, so
            # frozen() is happy while the mission is over. One window was spent
            # reporting "the counter never moved" about exactly that.
            if frozen.frozen(6.0)[0]:
                print("GUEST FROZEN — the counter was never going to move",
                      flush=True)
                return 3
            if not frozen.in_flight():
                print("NO LONGER IN FLIGHT — the mission ended; this window "
                      "proves nothing", flush=True)
                return 5
    nB = counter(w.fd)
    if nB == nA:
        print("counter never moved — no verification possible", flush=True)
        return 1
    vb = hud(f"{out}/b.png")
    print(f"counter {nA} -> {nB} (HUD {vb})", flush=True)

    nB2, namesB, wordsB = sample(w, defs, "B")
    cB = bit_counts(wordsB)
    nB_ents = len(namesB)
    surv_set = [(int(b), int(o)) for b, o in setA if cB[b, o] == nB2]
    surv_clr = [(int(b), int(o)) for b, o in clrA if cB[b, o] == nB_ents - nB2]
    print(f"[B] survivors — set-polarity {len(surv_set)}, "
          f"clear-polarity {len(surv_clr)}", flush=True)
    rows = ([{"polarity": "set", "offset": LO + o * 4, "bit": b} for b, o in surv_set]
            + [{"polarity": "clear", "offset": LO + o * 4, "bit": b} for b, o in surv_clr])
    json.dump({"nA": nA, "nB": nB2, "hud_a": v, "hud_b": vb,
               "entities_a": nA_ents, "entities_b": nB_ents,
               "candidates_a": {"set": len(setA), "clear": len(clrA)},
               "survivors": rows}, open(f"{out}/bitflag.json", "w"), indent=1)
    for r in rows[:40]:
        print(f"   {r['polarity']:5} pos{r['offset']:+#07x} bit {r['bit']:2d}",
              flush=True)
    if len(rows) > 40:
        print(f"   ... and {len(rows)-40} more (see bitflag.json)", flush=True)
    print(f"wrote {out}/bitflag.json", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())