#!/usr/bin/env python3
"""Phase 3: render a captured VPX FIFO stream to pixels.

Reconstructs the DPL scene graph from a VPX_FIFODUMP capture (see
decode_fifodump.py for the record format) and software-renders the frame each
vr_draw_scene commits, using the camera matrix delivered by the per-frame
camera action (31) and the view parameters from the view-node flush.

Rel4.10 wire protocol (established from the divrgb.scn capture, 2026-07-03):

  action 0   init            args string
  action 1   create          [type][name]         (host assigns node names)
  action 3   flush           [name][type][node struct fields]
  action 9   draw_scene      commit frame
  action 11  list_add        [parent][child]
  action 23  set_geom_verts  hdr [name][n_verts][3][n_blocks][1][5][n_verts][1.0f]
                             then float32 x,y,z per vertex (508-byte packetized)
  action 25  set_geom_conns  hdr [name][n_polys][verts_per_poly+1][0]
                             then indices, each poly a closed loop (last=first)
  action 31  camera          [?][view][3x3 rotation row-major][eye x,y,z]
  action 45  sync            token ping

  node types (from create/flush): 2=texture? 3=view 4=light 5=dcs 6=material(old)
    7=object 8=lod 9=geogroup 10=geometry 11=material
  type 9 geogroup flush: payload int 14 = material node name
  type 11 material flush: floats[10..12] = diffuse RGB
  type 3 view flush: floats include window (l,b,r,t), viewport w,h, near, far

Usage: render_capture.py <dump> [-o out.png] [--frame N]
"""
import struct
import sys

from PIL import Image, ImageDraw


def read_messages(path):
    msgs = []
    with open(path, "rb") as f:
        while True:
            hdr = f.read(8)
            if len(hdr) < 8:
                break
            if hdr[:4] != b"VPXM":
                raise SystemExit("bad magic")
            d = f.read(struct.unpack("<I", hdr[4:])[0])
            if len(d) >= 4:
                msgs.append((struct.unpack("<I", d[:4])[0], d[4:]))
    return msgs


def u32s(b):
    return list(struct.unpack(f"<{len(b) // 4}I", b[: len(b) // 4 * 4]))


def f32s(b):
    return list(struct.unpack(f"<{len(b) // 4}f", b[: len(b) // 4 * 4]))


class Scene:
    def __init__(self):
        self.types = {}        # name -> created type
        self.verts = {}        # geometry name -> [(x,y,z), ...]
        self.polys = {}        # geometry name -> [[i, ...], ...]
        self.material = {}     # material name -> (r, g, b)
        self.gg_material = {}  # geogroup name -> material name
        self.children = {}     # parent name -> [child names]
        self.view = None       # (win l,b,r,t, win-dist, vw, vh, near, far)
        self.background = (0, 0, 0)
        self.frames = []       # camera (3x3 rotation rows, eye) per draw_scene


def reconstruct(msgs):
    sc = Scene()
    camera = None
    geom_pend = None  # (name, n_verts) awaiting vertex float records
    conn_pend = None  # (name, n_polys, loop_len) awaiting index records
    for action, d in msgs:
        if action == 1:
            t, name = u32s(d)[:2]
            sc.types[name] = t
        elif action == 3:
            w = u32s(d)
            name, t = w[0], w[1]
            if t == 11 and len(d) >= 92:
                f = f32s(d)
                sc.material[name] = tuple(f[12:15])  # diffuse RGB
            elif t == 9 and len(d) >= 80:
                sc.gg_material[name] = w[16]
            elif t == 3 and len(d) >= 104:
                f = f32s(d)
                # window l,b,r,t; window-plane distance; viewport w,h;
                # near, far; background rgb  (last flush wins)
                sc.view = (f[6], f[7], f[8], f[9], f[10],
                           f[11], f[12], f[13], f[14])
                sc.background = tuple(f[15:18])
        elif action == 11:
            p, c = u32s(d)[:2]
            sc.children.setdefault(p, []).append(c)
        elif action == 23:
            if geom_pend is None:
                w = u32s(d)
                geom_pend = (w[0], w[2])
                sc.verts[w[0]] = []
            else:
                name, n = geom_pend
                f = f32s(d)
                vl = sc.verts[name]
                vl.extend((f[i], f[i + 1], f[i + 2])
                          for i in range(0, len(f) - 2, 3))
                if len(vl) >= n:
                    geom_pend = None
        elif action == 25:
            if conn_pend is None:
                w = u32s(d)
                conn_pend = (w[0], w[1], w[2])
                sc.polys[w[0]] = []
            else:
                name, n_polys, loop = conn_pend
                idx = u32s(d)
                pl = sc.polys[name]
                for i in range(0, len(idx), loop):
                    pl.append(idx[i:i + loop - 1])  # drop closing duplicate
                if len(pl) >= n_polys:
                    conn_pend = None
        elif action == 31:
            f = f32s(d)
            rot = [f[2:5], f[5:8], f[8:11]]
            eye = f[11:14]
            camera = (rot, eye)
        elif action == 9:
            sc.frames.append(camera)
    return sc


def render(sc, frame, out, ss=2):
    if sc.view:
        wl, wb, wr, wt, wd, vw, vh, near, far = sc.view
        vw, vh = int(vw), int(vh)
    else:
        wl, wb, wr, wt, wd = -1, -0.615, 1, 0.615, 1.3
        vw, vh, near, far = 832, 512, 2, 12000
    rot, eye = sc.frames[frame]
    W, H = vw * ss, vh * ss
    bg = tuple(max(0, min(255, int(c * 255 + 0.5))) for c in sc.background)
    img = Image.new("RGB", (W, H), bg)
    draw = ImageDraw.Draw(img)

    def project(p):
        x, y, z = (p[0] - eye[0], p[1] - eye[1], p[2] - eye[2])
        ex = rot[0][0] * x + rot[0][1] * y + rot[0][2] * z
        ey = rot[1][0] * x + rot[1][1] * y + rot[1][2] * z
        ez = rot[2][0] * x + rot[2][1] * y + rot[2][2] * z
        return ex, ey, ez

    # painter's algorithm: gather polys with depth, far first
    items = []
    for gg, mat in sc.gg_material.items():
        rgb = sc.material.get(mat, (1, 0, 1))
        col = tuple(max(0, min(255, int(c * 255 + 0.5))) for c in rgb)
        for geo in sc.children.get(gg, []):
            vl = sc.verts.get(geo)
            if not vl:
                continue
            evs = [project(v) for v in vl]
            for poly in sc.polys.get(geo, []):
                pts = [evs[i] for i in poly if i < len(evs)]
                if len(pts) < 3:
                    continue
                # view direction: Division looks down -Z in eye space
                depth = sum(p[2] for p in pts) / len(pts)
                if depth > -near:
                    continue  # behind the camera
                scr = []
                for ex, ey, ez in pts:
                    # Division screen x runs opposite to GL eye x: the SMPTE
                    # pattern (gray leftmost, -I/white/+Q PLUGE) comes out
                    # mirrored without the negation.
                    ndc_x = (-ex * wd / -ez - wl) / (wr - wl)
                    ndc_y = (ey * wd / -ez - wb) / (wt - wb)
                    scr.append((ndc_x * W, (1 - ndc_y) * H))
                items.append((depth, scr, col))
    items.sort(key=lambda it: it[0])
    for _, scr, col in items:
        draw.polygon(scr, fill=col)

    if ss > 1:
        img = img.resize((vw, vh), Image.LANCZOS)
    img.save(out)
    print(f"rendered frame {frame}: {len(items)} polys -> {out} ({vw}x{vh})")


def main():
    path = sys.argv[1]
    out = "capture.png"
    frame = 0
    if "-o" in sys.argv:
        out = sys.argv[sys.argv.index("-o") + 1]
    if "--frame" in sys.argv:
        frame = int(sys.argv[sys.argv.index("--frame") + 1])
    msgs = read_messages(path)
    sc = reconstruct(msgs)
    if "--eye" in sys.argv:  # override camera to survey the whole scene
        eye = [float(v) for v in sys.argv[sys.argv.index("--eye") + 1].split(",")]
        sc.frames = [([[1, 0, 0], [0, 1, 0], [0, 0, 1]], eye)]
        frame = 0
    print(f"{len(msgs)} messages: {len(sc.verts)} geometries, "
          f"{len(sc.material)} materials, {len(sc.frames)} frames, view={sc.view}")
    render(sc, frame, out)


if __name__ == "__main__":
    main()
