#!/usr/bin/env python3
"""Live world-state reader: the running game's entities, straight out of guest RAM.

Builds on gmem.py (Canary backs the guest address space with
/dev/shm/xenia_memory_*, so it is an ordinary file). The difference here is that
this is meant to run *in a loop* while the game plays, not against a snapshot:
the memory file is opened once and re-read with pread, and the expensive
whole-RAM vtable scan is done rarely and cached.

Two classes matter (see docs/re/structures/unit-struct-runtime.md):
  0x820af844  the parsed `.tbl` definition  — one per unit type
  0x820af030  the spawned entity instance   — one per thing in the scene

This module finds instances and reads their live transform. The transform
offsets are NOT guessed: `find_transform` scans an instance for a 3x3 block of
floats that is orthonormal to 1e-3 with det = +1, which a rotation matrix is and
essentially nothing else is.

Sub-commands:
    entities                     list live instances (name, address)
    probe <seconds> [hz] [out]   record every instance's raw bytes over time
    orient                       report orthonormal 3x3 blocks per instance
"""

import os
import re
import struct
import sys
import time

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

INST_VTABLE = 0x820AF030
DEF_VTABLE = 0x820AF844
WINDOW = 0x600          # bytes of an instance we read; its real size is unknown
NAME_STR_OFF = 0x10


class World:
    def __init__(self, path=None):
        self.path = path or gmem.mem_path()
        self.fd = os.open(self.path, os.O_RDONLY)
        self.size = os.path.getsize(self.path)
        self.instances = []          # [(va, file_off, name)]

    # ---------------------------------------------------------------- io
    def read(self, va, n):
        return os.pread(self.fd, n, gmem.va_to_off(va))

    def read_off(self, off, n):
        return os.pread(self.fd, n, off)

    def u32(self, va):
        return struct.unpack(">I", self.read(va, 4))[0]

    def cstr(self, va, n=96):
        b = os.pread(self.fd, n, gmem.va_to_off(va)).split(b"\0")[0]
        try:
            return b.decode("ascii")
        except UnicodeDecodeError:
            return None

    # ------------------------------------------------------------- scan
    def scan_vtable(self, vt):
        pat = struct.pack(">I", vt)
        hits = []
        for a, b in gmem.extents(self.fd, self.size):
            blob = os.pread(self.fd, b - a, a)
            start = 0
            while True:
                i = blob.find(pat, start)
                if i < 0:
                    break
                off = a + i
                if off % 4 == 0:
                    hits.append(off)
                start = i + 4
        return sorted(set(hits))

    def name_of(self, off):
        """The unit ID an object at `off` names, via its name-record pointer."""
        raw = self.read_off(off + 4, 4)
        if len(raw) < 4:
            return None
        (p,) = struct.unpack(">I", raw)
        try:
            return self.cstr(p + NAME_STR_OFF)
        except ValueError:
            return None

    def refresh(self, vt=INST_VTABLE):
        """Re-scan for live instances. Costs ~1 s; call it rarely."""
        out = []
        for off in self.scan_vtable(vt):
            nm = self.name_of(off)
            if nm and nm.startswith("UN_"):
                out.append((gmem.primary_va(off), off, nm))
        self.instances = out
        return out


# --------------------------------------------------------------- geometry


def floats(buf, off, n):
    return struct.unpack_from(">" + "f" * n, buf, off)


def is_rotation(m, tol=2e-3):
    """m = 9 floats, row-major. Orthonormal rows + det +1?"""
    r = [m[0:3], m[3:6], m[6:9]]
    for row in r:
        n2 = sum(c * c for c in row)
        if abs(n2 - 1.0) > tol:
            return False
    for i in range(3):
        for j in range(i + 1, 3):
            if abs(sum(r[i][k] * r[j][k] for k in range(3))) > tol:
                return False
    det = (r[0][0] * (r[1][1] * r[2][2] - r[1][2] * r[2][1])
           - r[0][1] * (r[1][0] * r[2][2] - r[1][2] * r[2][0])
           + r[0][2] * (r[1][0] * r[2][1] - r[1][1] * r[2][0]))
    return abs(det - 1.0) < 1e-2


def find_rotations(buf, stride=4):
    """Offsets where 9 consecutive floats form a rotation matrix.

    Also tries the 4-float-per-row (3x4 / 4x4 matrix) stride, which is how a
    transform with a translation column is normally stored.
    """
    out = []
    for off in range(0, len(buf) - 36, stride):
        m = floats(buf, off, 9)
        if all(abs(c) <= 1.001 for c in m) and is_rotation(m):
            out.append((off, "3x3", m))
    for off in range(0, len(buf) - 48, stride):
        rows = [floats(buf, off + 16 * i, 3) for i in range(3)]
        m = rows[0] + rows[1] + rows[2]
        if all(abs(c) <= 1.001 for c in m) and is_rotation(m):
            out.append((off, "4x4", m))
    return out


# ------------------------------------------------------------------ cmds


def cmd_entities(w, args):
    t0 = time.time()
    inst = w.refresh()
    print(f"# {len(inst)} live instances (scan {time.time()-t0:.2f}s)")
    from collections import Counter
    c = Counter(n for _, _, n in inst)
    for nm, k in c.most_common():
        print(f"  {k:4d}  {nm}")


def cmd_orient(w, args):
    inst = w.refresh()
    want = args[0] if args else None
    for va, off, nm in inst:
        if want and want not in nm:
            continue
        buf = w.read_off(off, WINDOW)
        rots = find_rotations(buf)
        print(f"\n{va:#010x} {nm}  -> {len(rots)} rotation-like block(s)")
        for o, kind, m in rots[:6]:
            print(f"   +{o:#05x} {kind}  "
                  + " ".join(f"{v:+.3f}" for v in m))


def cmd_probe(w, args):
    secs = float(args[0]) if args else 6.0
    hz = float(args[1]) if len(args) > 1 else 5.0
    out = args[2] if len(args) > 2 else "probe.bin"
    inst = w.refresh()
    print(f"# probing {len(inst)} instances for {secs}s at {hz}Hz -> {out}")
    with open(out, "wb") as f:
        f.write(struct.pack("<II", len(inst), WINDOW))
        for va, off, nm in inst:
            nb = nm.encode()[:63].ljust(64, b"\0")
            f.write(struct.pack("<II", va, off) + nb)
        n = int(secs * hz)
        for i in range(n):
            t = time.time()
            f.write(struct.pack("<d", t))
            for va, off, nm in inst:
                f.write(w.read_off(off, WINDOW).ljust(WINDOW, b"\0"))
            time.sleep(max(0, 1.0 / hz - (time.time() - t)))
    print(f"# wrote {n} frames")


def main():
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    w = World()
    cmd, args = sys.argv[1], sys.argv[2:]
    {"entities": cmd_entities, "orient": cmd_orient, "probe": cmd_probe}[cmd](w, args)


if __name__ == "__main__":
    main()