#!/usr/bin/env python3
"""Does an `OB`-badged entity carry a flag, and is `REMAINING OB` its count?

`mission-objective-counter.md` has the counter's address; what it *counts* is the
part the autopilot needs, because "shoot what closes the mission" requires
picking the right target, not knowing how many are left.

Two questions, in order of how cheaply they can be killed:

1. **Is the counter just a per-class head-count?** Print the class histogram
   beside the counter. The corpus already suspects not (012 on the HUD against
   118 live ADAN), and one run settles it for every class at once.
2. **Is there a per-entity flag whose set-cardinality is the counter?** For every
   4-byte offset in a window around each entity, count how many entities share
   each value. An offset where exactly N entities agree, with N the counter, is a
   candidate — and there will be many by chance, so the answer is the SECOND
   sample: after the counter moves to N', the same (offset, value) must be shared
   by exactly N' entities. That is the same "verify across a transition you did
   not select on" rule the address itself had to pass.

🔴 Known limit, stated because it bounds the conclusion: `entities2.typed`
enumerates entities by their position triple CHANGING between two samples, so a
stationary objective is invisible to it. Stage 02's objective is "shoot down all
invading enemy fighters", which move — but a null result here does not rule out a
flag on objects this enumeration never sees.

Usage: ob_flag.py <out-dir> [timeout_s]
"""
import json
import os
import struct
import subprocess
import sys
import time
from collections import Counter, defaultdict

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 is NOT at a fixed address across runs — see
# structures/mission-objective-counter.md. 0xbdb59668 is the value it takes in
# 3 of the 5 runs measured; when it is wrong this script refuses to run rather
# than reporting nonsense, and ob_hunt.py finds the run's own address in ~1 min
# of flight. Override with OB_VA.
VA = int(os.environ.get("OB_VA", "0xBDB59668"), 0)
RADIUS = 0x400
DELTA = 0x130


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


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, out, tag):
    """(counter, entity list, {(offset, value): [entity positions]})."""
    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())
    groups = defaultdict(list)
    for off, nm in ents:
        lo = off - RADIUS
        blob = os.pread(fd, RADIUS * 2, lo)
        for k in range(0, len(blob) - 3, 4):
            groups[(lo + k - off, blob[k:k + 4])].append(nm)
    print(f"[{tag}] counter={n} entities={len(ents)}", flush=True)
    return n, ents, groups


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

    # RETRY an unreadable frame rather than counting it as a mismatch. ob_read
    # returns None on a frame it cannot read - an explosion across the plate, a
    # flash, the HUD momentarily gone - and treating that as "the address is
    # wrong" aborted a perfectly good run once.
    v = None
    for _ in range(8):
        v = hud(f"{out}/a.png")
        n0 = counter(w.fd)
        if v is not None:
            break
        time.sleep(3)
    print(f"HUD={v} RAM={n0}", flush=True)
    if v is None:
        print("could not read the HUD counter at all in 8 tries — is this "
              "even in flight?", flush=True)
        return 2
    if v != n0:
        print("HUD and RAM disagree — wrong address for this run; re-scan with "
              "ob_hunt.py before trusting anything below", flush=True)
        return 2

    # RETRY an empty sample. entities2.moving() types entities by their position
    # CHANGING between two reads, so a sample that lands on a load, a stall, or
    # simply loses the race with another scanner comes back with nothing — and a
    # zero-entity sample silently produces zero candidates and a void run. Seen
    # once, with sample B finding 81 entities moments later.
    for _ in range(6):
        nA, entsA, gA = sample(w, defs, out, "A")
        if len(entsA) >= 10:
            break
        print("[A] too few entities — retrying", flush=True)
        time.sleep(5)
    if len(entsA) < 10:
        print("[A] never got a usable entity sample", flush=True)
        return 4
    hist = Counter(nm for _, nm in entsA)
    print(f"[A] class histogram vs counter {nA}:", flush=True)
    for nm, k in hist.most_common(12):
        print(f"      {k:4d}  {nm}", flush=True)
    exact = [nm for nm, k in hist.items() if k == nA]
    print(f"[A] classes whose head-count equals the counter: {exact or 'NONE'}",
          flush=True)

    candA = {k: v for k, v in gA.items() if len(v) == nA}
    print(f"[A] offsets where exactly {nA} entities agree: {len(candA)}", flush=True)
    # WHO a candidate groups is a discriminator available without waiting for a
    # transition: an OB flag should mark a subset of the hostiles, not a mix that
    # includes the player and its wingmen. Printed for every candidate rather
    # than filtered on, because "objectives are always hostile" is an assumption
    # about the mission, not a measurement -- an escort objective would be
    # friendly, and Stage 02 has one (the ACROPOLIS).
    for (d, val), mem in sorted(candA.items()):
        who = Counter(nm.replace("UN_", "") for nm in mem)
        print(f"      pos{d:+#07x} = {val.hex()}  {dict(who)}", flush=True)

    stuck = 0
    while time.time() < deadline:
        time.sleep(5)
        if counter(w.fd) != nA:
            break
        stuck += 1
        if stuck % 12 == 0:
            if frozen.frozen(6.0)[0]:
                print("GUEST FROZEN — the world stopped advancing, so the "
                      "counter was never going to move; this run proves "
                      "nothing", 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)

    _, entsB, gB = sample(w, defs, out, "B")
    survivors = {k: (len(gA[k]), len(gB.get(k, []))) for k in candA
                 if len(gB.get(k, [])) == nB}
    print(f"[B] of {len(candA)} candidates, {len(survivors)} still hold "
          f"exactly {nB}", flush=True)
    rows = [{"delta": d, "value": val.hex(), "a": a, "b": b}
            for (d, val), (a, b) in sorted(survivors.items())]
    json.dump({"nA": nA, "nB": nB, "hud_a": v, "hud_b": vb,
               "entities_a": len(entsA), "entities_b": len(entsB),
               "class_matches": exact, "candidates_a": len(candA),
               "survivors": rows}, open(f"{out}/flag.json", "w"), indent=1)
    for r in rows[:40]:
        print(f"   pos{r['delta']:+#07x} = {r['value']}  {r['a']} -> {r['b']}",
              flush=True)
    print(f"wrote {out}/flag.json", flush=True)
    return 0


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