"""M4c: the authentic backend -- consumes the vpxlog device's VPXM wire stream
(offline fifodump file, or live VPX_FIFOSOCK) and renders through emu860c +
the GPU tile path. Frames written as live_NNNN.png.

  python live_server.py <capture.fifodump> [max_frames]
  python live_server.py sock:<port>        [max_frames]   (live)
"""
import sys, os, time, struct, socket
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
SRC = sys.argv[1]
MAXF = int(sys.argv[2]) if len(sys.argv) > 2 else 8
W, H = 832, 512


class Source:
    """Incremental VPXM record source: file or socket."""
    def __init__(self, spec):
        self.buf = b''
        self.sock = None
        if spec.startswith('sock:'):
            port = int(spec[5:])
            self.sock = socket.socket()
            while True:
                try:
                    self.sock.connect(('127.0.0.1', port))
                    break
                except OSError:
                    time.sleep(0.5)
            self.sock.settimeout(0.5)
            print("connected to fifosock :%d" % port, flush=True)
        else:
            self.buf = open(spec, 'rb').read()
            print("loaded %d bytes from %s" % (len(self.buf), spec), flush=True)
        self.off = 0

    def poll(self):
        if self.sock:
            try:
                data = self.sock.recv(1 << 16)
                if data:
                    self.buf += data
            except socket.timeout:
                pass

    def next_record(self):
        while True:
            i = self.buf.find(b'VPXM', self.off)
            if i < 0 or i + 8 > len(self.buf):
                return None
            ln = struct.unpack_from('<I', self.buf, i + 4)[0]
            if i + 8 + ln > len(self.buf):
                return None
            body = self.buf[i + 8:i + 8 + ln]
            self.off = i + 8 + ln
            if len(body) >= 4:
                action = struct.unpack_from('<I', body, 0)[0]
                if action < 0x100:
                    return (action, body[4:])
            # skip non-message bursts


src = Source(SRC)
# preload the whole wire as the live queue (offline); production firmware
_recs = []
while True:
    rec = src.next_record()
    if rec is None:
        src.poll()
        if src.sock is None:
            break
        continue
    _recs.append(rec)
    if src.sock is None and len(_recs) > 200000:
        break
print("queued %d records" % len(_recs), flush=True)
r = boot(fw='vrend410', queue=_recs)
shim = CpuShim()
emu860c.watch_add(0x08020000, 0x08190000)
g = igc_gpu.GpuTile()
print("GPU:", g.ctx.info['GL_RENDERER'], flush=True)

texu_acc = np.zeros((H, W), np.int32)
frames_done = 0
prev_was_draw = False
t0 = time.time()


def render_frame(nth):
    raw = emu860c.watch_drain()
    vals = np.frombuffer(raw, dtype=np.uint32).reshape(-1, 2)[:, 1] \
        if len(raw) else np.zeros(0, np.uint32)
    sends, order = [], []
    i = 0
    while i + 1 < len(vals):
        a, op = int(vals[i]), int(vals[i + 1])
        c = (op >> 28) & 0xf
        if c in (1, 9) and 0x08000000 <= a < 0x08020000:
            sends.append([emu860c.r32(a + 4 * k) for k in range(op & 0x7f)])
        elif c == 2:
            order.append((a, len(sends)))
        i += 2
    if not order:
        return False
    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_acc[gy, gx] = igc_gpu.GpuTile.rdbits(row[px], 57, 20)
        ntiles += 1
    if not ntiles:
        return False
    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_acc[:, 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, 'live_%04d.png' % nth))
    print("frame %d: %d tiles, %d sends (cmd %d)" %
          (nth, ntiles, len(sends), r.qi), flush=True)
    return True


RECEIVE_PC = emu_main.MAPS['vrend410']['receive']
idle = 0
last_note = time.time()
while frames_done < MAXF and idle < 40:
    reason, steps = emu860c.run(500_000_000)
    if time.time() - last_note > 20:
        last_note = time.time()
        st = emu860c.getstate()
        print("... cmd %d, %.1fB steps, pc=%#x" %
              (r.qi, st['steps'] / 1e9, st['pc']), flush=True)
    if reason == 3:
        continue                      # long processing slice; keep going
    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 == RECEIVE_PC:
        if prev_was_draw:
            if render_frame(frames_done):
                frames_done += 1
            prev_was_draw = False
        if r.qi >= len(r.queue):
            break
        nxt = r.queue[r.qi][0]
        if ANAME.get(nxt) == 'draw_scene':
            emu860c.watch_drain()
            prev_was_draw = True
    if h(shim) == 'done':
        break
print("done: %d frames in %.1fs (cmds %d)" %
      (frames_done, time.time() - t0, r.qi), flush=True)