#!/usr/bin/env python3
"""Dump the PAINT ORDER of every live UI screen object, off a running Canary.

A screen object (vtable `0x820b30b4`) holds its elements at `+0x08` as
`{ptr, count, capacity}` over 48-byte records in **declaration** order, and a
second list at `+0x30` — pointers to the same records, in the order the screen
is **painted**. `docs/re/structures/ui-screen-runtime.md` establishes both.

This walks every resident screen object and prints, per object, the permutation
and each element's pivot, which is what identifies the build in the file (a
declaration entry's pivot is exactly half the decoded sprite).

    screen_children.py [/dev/shm/xenia_memory_XXX]

Point `$GMEM_FILE` at a snapshot to work offline:
    cp --sparse=always /dev/shm/xenia_memory_* /tmp/snap.bin
"""
import os
import struct
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gmem

VTABLE = 0x820B30B4
ELEMENTS_AT = 0x08   # {ptr, count, capacity} -> 48-byte element records
CHILDREN_AT = 0x30   # {ptr, count, capacity} -> pointers into those records
RECORD = 48


def main():
    path = gmem.mem_path()
    size = os.path.getsize(path)
    with open(path, "rb") as f:
        def rd(va, n):
            f.seek(gmem.va_to_off(va))
            return f.read(n)

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

        def f32(va):
            return struct.unpack(">f", rd(va, 4))[0]

        # every object whose word 0 is the screen vtable
        objs = []
        pat = struct.pack(">I", VTABLE)
        CHUNK = 1 << 24
        for start, end in gmem.extents(f.fileno(), size):
            pos = start
            while pos < end:
                f.seek(pos)
                buf = f.read(min(CHUNK, end - pos))
                if not buf:
                    break
                i = buf.find(pat)
                while i != -1:
                    va = gmem.primary_va(pos + i)
                    if va is not None:
                        objs.append(va)
                    i = buf.find(pat, i + 1)
                pos += len(buf)

        print(f"{len(objs)} object(s) with vtable {VTABLE:#010x}")
        for obj in objs:
            try:
                eptr, ecount = u32(obj + ELEMENTS_AT), u32(obj + ELEMENTS_AT + 4)
                cptr, ccount = u32(obj + CHILDREN_AT), u32(obj + CHILDREN_AT + 4)
            except ValueError:
                continue
            if not (0 < ecount <= 4096) or eptr == 0:
                continue
            # element record pointer -> declaration index
            index_of, pivots = {}, []
            ok = True
            for i in range(ecount):
                rec = eptr + i * RECORD
                try:
                    index_of[u32(rec)] = i
                    pivots.append((f32(rec + 0x10), f32(rec + 0x14)))
                except ValueError:
                    ok = False
                    break
            if not ok:
                continue
            print(f"\n== object {obj:#010x}: {ecount} elements, {ccount} children")
            print("   pivots: " + " ".join(f"{i}:({x:g},{y:g})" for i, (x, y) in enumerate(pivots)))
            if not (0 < ccount <= 4096) or cptr == 0:
                print("   (no child array)")
                continue
            order, unknown = [], 0
            for i in range(ccount):
                p = u32(cptr + i * 4)
                if p in index_of:
                    order.append(index_of[p])
                else:
                    order.append(None)
                    unknown += 1
            print(f"   paint order: {order}")
            if unknown:
                print(f"   ({unknown} child pointer(s) did not match an element record)")


if __name__ == "__main__":
    main()