"""Battle content frames: production firmware on the C core, real game wire,
per-draw EFFECT-PRIMITIVE readout (render_fx path) -> actual scene geometry.
The draws build effect programs at ~0x815f000 (like fxtest's 0x816f000)."""
import sys, os, time, struct
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE); sys.path.insert(0, os.path.dirname(HERE))
sys.path.insert(0, r'C:\VWE\TeslaRel410\dpl3-revive\patha')
import emu860, emu_main, emu860c, igc_exec, igc_gpu, igc_gpu_frame
import numpy as np
from PIL import Image
from driver import boot, CpuShim
from vrboard import A
ANAME = {int(a): a.name for a in A}
emu860.Mem.log = lambda self, *a, **k: None

SRC = sys.argv[1]
MAXF = int(sys.argv[2]) if len(sys.argv) > 2 else 8
PROG_LO, PROG_HI = 0x08158000, 0x08170000    # per-draw effect program window
W, H = 832, 512

data = open(SRC, 'rb').read()
recs = []; off = 0
while True:
    i = data.find(b'VPXM', off)
    if i < 0: break
    ln = struct.unpack_from('<I', data, i + 4)[0]
    body = data[i + 8:i + 8 + ln]; off = i + 8 + ln
    if len(body) >= 4:
        a = struct.unpack_from('<I', body, 0)[0]
        if a < 0x100: recs.append((a, body[4:]))
print("queued %d records" % len(recs), flush=True)

r = boot(fw='vrend410', queue=recs)
shim = CpuShim()
RECV = emu_main.MAPS['vrend410']['receive']
ctx = igc_gpu.GpuTile().ctx
print("GPU:", ctx.info['GL_RENDERER'], flush=True)


def f32(w): return struct.unpack('<f', struct.pack('<I', w & 0xffffffff))[0]


def render_effects(nth):
    prog = {a: emu860c.r32(a) for a in range(PROG_LO, PROG_HI, 4)
            if emu860c.r32(a)}
    def rd(a): return prog.get(a, 0)
    setups = [a for a in sorted(prog)
              if rd(a) == 0x100 and ((rd(a + 4) >> 8) & 0xff) == 0x2c]
    if not setups:
        return False, 0
    recs4 = []
    for a in setups:
        edges = []
        p = a + 4
        while ((rd(p) >> 8) & 0xff) == 0x2c and len(edges) < 6:
            edges.append((f32(rd(p+4)), f32(rd(p+8)), f32(rd(p+12)))); p += 16
        zp = up = vp = None; q = p
        for _ in range(40):
            wv = rd(q); op = (wv >> 8) & 0xff; ad = wv & 0xff
            if op == 0x21 and zp is None: zp = (f32(rd(q+4)), f32(rd(q+8)), f32(rd(q+12)))
            if op == 0x43 and ad == 58 and up is None: up = (f32(rd(q+4)), f32(rd(q+8)), f32(rd(q+12)))
            if op == 0x43 and ad == 78 and vp is None: vp = (f32(rd(q+4)), f32(rd(q+8)), f32(rd(q+12)))
            if zp and up and vp: break
            q += 4
        if zp is None or len(edges) < 2:
            continue
        hasuv = 1.0 if (up and vp) else 0.0
        up = up or (0, 0, 0); vp = vp or (0, 0, 0)
        eslots = [(e[0], e[1], e[2], 1.0) for e in edges[:4]]
        while len(eslots) < 4: eslots.append((0.0, 0.0, 0.0, 0.0))
        recs4 += eslots + [(*zp, float(len(edges))), (*up, hasuv), (*vp, 0.0)]
    if not recs4:
        return False, 0
    prims = np.array(recs4, dtype=np.float32).reshape(-1, 4)
    n_prims = len(prims) // 7
    prog_gl = ctx.compute_shader(igc_gpu_frame.FX_SHADER)
    b0 = ctx.buffer(prims.tobytes()); out = ctx.buffer(reserve=W * H * 16)
    b0.bind_to_storage_buffer(0); out.bind_to_storage_buffer(1)
    prog_gl['n_prims'] = n_prims; prog_gl['width'] = W
    prog_gl.run(group_x=(W + 63) // 64, group_y=H)
    raw = np.frombuffer(out.read(), dtype=np.uint32).reshape(H, W, 4)
    z = raw[..., 0].view(np.float32); u = raw[..., 1].view(np.float32); v = raw[..., 2].view(np.float32)
    m = z > -1e29
    rgb = np.zeros((H, W, 3), np.uint8)
    if m.any():
        un = np.zeros_like(u); vn = np.zeros_like(v)
        un[m] = (u[m] - u[m].min()) / max(1e-9, np.ptp(u[m]))
        vn[m] = (v[m] - v[m].min()) / max(1e-9, np.ptp(v[m]))
        rgb[..., 0][m] = (60 + 180 * un[m]).astype(np.uint8)
        rgb[..., 1][m] = (40 + 160 * vn[m]).astype(np.uint8)
        rgb[..., 2][m] = (220 - 140 * un[m]).astype(np.uint8)
    b0.release(); out.release()
    Image.fromarray(rgb, 'RGB').save(os.path.join(HERE, 'battle_%04d.png' % nth))
    return True, (n_prims, int(m.sum()))


frames = 0; prev_draw = False; t0 = time.time()
while frames < MAXF and r.qi < len(r.queue):
    reason, _ = emu860c.run(500_000_000)
    if reason == 3: continue
    if reason != 0:
        print("core reason %d" % reason, flush=True); break
    pc = emu860c.getstate()['pc']
    h = r.hooks.get(pc)
    if h is None:
        print("sentinel", flush=True); break
    if pc == RECV:
        if prev_draw:
            ok, info = render_effects(frames)
            if ok:
                print("battle frame %d: %s prims/px (cmd %d)" % (frames, info, r.qi), flush=True)
                frames += 1
            prev_draw = False
        if r.qi < len(r.queue) and ANAME.get(r.queue[r.qi][0]) == 'draw_scene':
            prev_draw = True
    if h(shim) == 'done':
        break
print("done: %d battle frames in %.1fs (cmd %d)" % (frames, time.time() - t0, r.qi), flush=True)