#!/usr/bin/env python3
"""Phase 3d: render a captured *game* VPX FIFO stream (full DPL hierarchy).

Unlike flyk's flat DIVRGB scene (render_capture.py), the game builds the real
DPL graph: instances reference objects and DCS transform nodes; objects hold
LODs holding geogroups holding geometries; dcs_link builds an articulation
tree of 4x4 matrices (mech torso/legs/arms). This reconstructs that graph
from a VPX_FIFODUMP capture and renders the scene each draw_scene commits.

Node types (empirical, Rel4.10 wire): 2=? 3=view 4=instance 5=dcs 6=material
(old) 7=object 8=lod 9=geogroup 10=geometry 11=material 12=texmap 13=? 14=?

Usage: render_game.py <dump> [-o out.png] [--frame N] [--eye x,y,z]
"""
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 U(b):
    return list(struct.unpack(f"<{len(b) // 4}I", b[: len(b) // 4 * 4]))


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


def mat_mul(a, b):
    """4x4 (row-major, row-vector convention: v' = v @ M)"""
    return [[sum(a[i][k] * b[k][j] for k in range(4)) for j in range(4)]
            for i in range(4)]


def mat_id():
    return [[1.0 if i == j else 0.0 for j in range(4)] for i in range(4)]


def xform(v, m):
    x, y, z = v
    return (x * m[0][0] + y * m[1][0] + z * m[2][0] + m[3][0],
            x * m[0][1] + y * m[1][1] + z * m[2][1] + m[3][1],
            x * m[0][2] + y * m[1][2] + z * m[2][2] + m[3][2])


class Scene:
    def __init__(self):
        self.types = {}
        self.verts = {}
        self.polys = {}
        self.material = {}
        self.gg_material = {}
        self.children = {}     # list_add: parent -> [children]
        self.inst_object = {}  # instance -> object
        self.inst_dcs = {}     # instance -> dcs
        self.dcs_mat = {}      # dcs -> 4x4 local matrix
        self.dcs_parent = {}   # dcs child -> parent (dcs_link)
        self.view = None
        self.background = (0, 0, 0)
        self.frames = []


def parse_dcs_matrix(f):
    """The 132-byte dcs flush payload: [name][type] then fields; the 4x4 is
    the 16 floats starting at float index 4 (rows [x,y,z,0], row 3 = T)."""
    m = [f[4 + r * 4: 8 + r * 4] for r in range(4)]
    # sanity: last column should be ~(0,0,0,1)
    if abs(m[3][3] - 1.0) > 0.5 and abs(m[0][3]) < 0.01:
        # some builds put w=1 elsewhere; force affine
        m[3][3] = 1.0
    for r in range(3):
        m[r][3] = 0.0
    m[3][3] = 1.0
    return m


def reconstruct(msgs):
    sc = Scene()
    camera = None
    geom_pend = None
    conn_pend = None
    for action, d in msgs:
        if action == 1:
            w = U(d)
            if len(w) >= 2:
                sc.types[w[1]] = w[0]
        elif action == 3:
            w = U(d)
            if len(w) < 2:
                continue
            name, t = w[0], w[1]
            if t == 11 and len(d) >= 92:
                f = F(d)
                sc.material[name] = tuple(f[12:15])
            elif t == 9 and len(d) >= 80:
                sc.gg_material[name] = w[16]
            elif t == 3 and len(d) >= 104:
                f = F(d)
                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 t == 5 and len(d) >= 132:
                sc.dcs_mat[name] = parse_dcs_matrix(F(d))
            elif t == 4:
                # instance: find object/dcs refs by node-type lookup
                for val in w[2:]:
                    if val and val != 0xFFFFFFFF:
                        vt = sc.types.get(val)
                        if vt == 7:
                            sc.inst_object[name] = val
                        elif vt == 5:
                            sc.inst_dcs[name] = val
        elif action == 7 and len(d) >= 8:   # dcs_link parent -> child
            w = U(d)
            sc.dcs_parent[w[1]] = w[0]
        elif action == 11 and len(d) >= 8:
            w = U(d)
            sc.children.setdefault(w[0], []).append(w[1])
        elif action == 23:
            if geom_pend is None:
                w = U(d)
                if len(d) >= 36:
                    # header: [name][0][n_verts][stride_floats][n_msgs]...
                    geom_pend = [w[0], w[2], w[3], []]
                    sc.verts[w[0]] = []
            else:
                name, n, stride, acc = geom_pend
                acc.extend(F(d))
                if len(acc) >= n * stride:
                    sc.verts[name] = [
                        (acc[i], acc[i + 1], acc[i + 2])
                        for i in range(0, n * stride, stride)]
                    geom_pend = None
        elif action == 25:
            if conn_pend is None:
                w = U(d)
                if len(d) >= 16:
                    conn_pend = (w[0], w[1], w[2])
                    sc.polys[w[0]] = []
            else:
                name, n_polys, loop = conn_pend
                idx = U(d)
                pl = sc.polys[name]
                if loop >= 2:
                    for i in range(0, len(idx), loop):
                        pl.append(idx[i:i + loop - 1])
                if len(pl) >= n_polys:
                    conn_pend = None
        elif action == 31:
            f = F(d)
            camera = ([f[2:5], f[5:8], f[8:11]], f[11:14])
        elif action == 9:
            sc.frames.append(camera)
    return sc


def dcs_world(sc, dcs, cache, depth=0):
    if dcs in cache:
        return cache[dcs]
    m = sc.dcs_mat.get(dcs, mat_id())
    p = sc.dcs_parent.get(dcs)
    if p is not None and p != dcs and depth < 64:
        m = mat_mul(m, dcs_world(sc, p, cache, depth + 1))
    cache[dcs] = m
    return m


def gather_polys(sc):
    """instance -> object -> lod -> geogroup -> geometry, transformed.

    Placement: instances are list_add children of DCS nodes (dcs -> instance);
    dcs_link builds the dcs->dcs articulation tree above them."""
    inst_parent = {}
    for parent, kids in sc.children.items():
        if sc.types.get(parent) == 5:
            for k in kids:
                if sc.types.get(k) == 4:
                    inst_parent[k] = parent
    out = []
    cache = {}
    for inst, obj in sc.inst_object.items():
        world = mat_id()
        d = inst_parent.get(inst, sc.inst_dcs.get(inst))
        if d is not None:
            world = dcs_world(sc, d, cache)
        for lod in sc.children.get(obj, []):
            # use only the first (highest-detail) LOD child set
            ggs = sc.children.get(lod, [])
            if not ggs:
                continue
            for gg in ggs:
                rgb = sc.material.get(sc.gg_material.get(gg, -1),
                                      (0.8, 0.2, 0.8))
                col = tuple(max(0, min(255, int(c * 255 + .5))) for c in rgb)
                for geo in sc.children.get(gg, []):
                    vl = sc.verts.get(geo)
                    if not vl:
                        continue
                    wv = [xform(v, world) for v in vl]
                    for poly in sc.polys.get(geo, []):
                        pts = [wv[i] for i in poly if i < len(wv)]
                        if len(pts) >= 3:
                            out.append((pts, col))
            break  # first LOD only
    return out


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.6154, 1, 0.6154, 1.732
        vw, vh, near, far = 832, 512, 0.25, 1150
    cam = sc.frames[frame]
    if cam is None:
        raise SystemExit("no camera at that frame")
    rot, eye = cam
    W, H = vw * ss, vh * ss
    bg = tuple(max(0, min(255, int(c * 255 + .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

    def clip_near(evs, zn):
        """Sutherland-Hodgman clip against eye-space plane z = -zn."""
        out = []
        n = len(evs)
        for i in range(n):
            a, b = evs[i], evs[(i + 1) % n]
            ain, bin_ = a[2] <= -zn, b[2] <= -zn
            if ain:
                out.append(a)
            if ain != bin_:
                t = (-zn - a[2]) / (b[2] - a[2])
                out.append((a[0] + t * (b[0] - a[0]),
                            a[1] + t * (b[1] - a[1]), -zn))
        return out

    polys = gather_polys(sc)
    print(f"gathered {len(polys)} world polys")
    items = []
    for pts, col in polys:
        evs = clip_near([project(p) for p in pts], near)
        if len(evs) < 3:
            continue
        depth = sum(p[2] for p in evs) / len(evs)
        scr = []
        for ex, ey, ez in evs:
            # game world is y-down (DCS matrices carry a y reflection), so
            # screen-up = -eye_y; screen-x mirrored as with flyk.
            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 = "game_frame.png"
    frame = -1
    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)
    print(f"{len(msgs)} msgs: {len(sc.verts)} geoms, {len(sc.material)} mats, "
          f"{len(sc.inst_object)} placed instances, {len(sc.dcs_mat)} dcs, "
          f"{len(sc.frames)} frames, view={sc.view}")
    if "--eye" in sys.argv:
        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
    render(sc, frame, out)


if __name__ == "__main__":
    main()