"""M4b: offline end-to-end frames -- run cap7 through the C core with a
write-watch on the DMA queue pages; at each draw boundary, reconstruct the
frame's DMA stream + payloads and render the tiles on the GPU. Writes
frame_NNNN.png for the first N draws."""
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
import igc_exec, igc_gpu
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
NFRAMES = int(sys.argv[1]) if len(sys.argv) > 1 else 10
W, H = 832, 512

r = boot()
shim = CpuShim()
# watch the DMA queue-page region (the coeff-copy writes land here)
emu860c.watch_add(0x08020000, 0x08190000)
g = igc_gpu.GpuTile()
print("GPU:", g.ctx.info['GL_RENDERER'], flush=True)

frames_done = 0
prev_was_draw = [False]
t0 = time.time()
texu_acc = np.zeros((H, W), np.int32)     # persistent across draws (fields)


def render_frame(nth):
    raw = emu860c.watch_drain()
    words = np.frombuffer(raw, dtype=np.uint32).reshape(-1, 2)
    vals = words[:, 1]
    # DMA pairs: consecutive (addr, op) values in write order
    tiles = {}
    i = 0
    n = len(vals)
    while i + 1 < n:
        a, op = int(vals[i]), int(vals[i + 1])
        c = (op >> 28) & 0xf
        if c in (1, 9) and 0x08000000 <= a < 0x08020000:
            sz = op & 0x7f
            payload = [emu860c.r32(a + 4 * k) for k in range(sz)]
            tiles.setdefault('sends', []).append(payload)
        elif c == 2:
            tid = a
            tiles.setdefault('order', []).append((tid, len(tiles.get('sends', []))))
        i += 2
    sends = tiles.get('sends', [])
    order = tiles.get('order', [])
    if not order:
        return False
    texu = texu_acc
    prev_cut = 0
    ntiles = 0
    for tid, cut in order:
        ox = (tid & 0x1f) * 64
        oy = ((tid >> 5) & 0x1f) * 128
        instrs = []
        for pl in sends[prev_cut:cut]:
            ins, _ = igc_exec.parse(pl)
            instrs += ins
        prev_cut = cut
        if ox >= W or oy >= H or not instrs:
            continue
        data = g.run(instrs, ox=ox, oy=oy, pre_seed_texz_x=True)
        for py_ in range(igc_gpu.TILE_H):
            gy = oy + py_
            if gy >= H:
                break
            row = data[py_ * igc_gpu.TILE_W:(py_ + 1) * igc_gpu.TILE_W]
            for px in range(igc_gpu.TILE_W):
                gx = ox + px
                if gx >= W:
                    break
                texu[gy, gx] = igc_gpu.GpuTile.rdbits(row[px], 57, 20)
        ntiles += 1
    PAL = [(180, 180, 180), (180, 180, 16), (16, 180, 180), (16, 180, 16),
           (180, 16, 180), (180, 16, 16), (16, 16, 180)]
    rgb = np.zeros((H, W, 3), np.uint8)
    for gx in range(W):
        u = int(np.median(texu[:, gx]))
        rgb[:, gx] = (0, 0, 0) if u < 48 else PAL[min(6, (u - 48) * 7 // (W - 48))]
    Image.fromarray(rgb, 'RGB').save(os.path.join(HERE, 'frame_%04d.png' % nth))
    print("frame %d: %d tiles, %d sends" % (nth, ntiles, len(sends)), flush=True)
    return True


while frames_done < NFRAMES:
    reason, steps = emu860c.run(2_000_000_000)
    if reason != 0:
        print("reason %d, stopping" % reason, flush=True)
        break
    pc = emu860c.getstate()['pc']
    h = r.hooks.get(pc)
    if h is None:
        break
    # draw boundaries are receive->receive spans: render only when we are
    # BACK at h_receive after a draw_scene delivery (igcwait etc. fire mid-draw)
    if pc == 0xf04024c0:      # h_receive
        if prev_was_draw[0]:
            if render_frame(frames_done):
                frames_done += 1
            prev_was_draw[0] = False
        nxt = r.qi
        if nxt < len(r.queue) and ANAME.get(r.queue[nxt][0]) == 'draw_scene':
            emu860c.watch_drain()      # discard pre-draw noise
            prev_was_draw[0] = True
    if h(shim) == 'done':
        break
print("done: %d frames in %.1fs (cmd %d)" % (frames_done, time.time() - t0, r.qi), flush=True)