"""gpu_raster.py -- the verified real-time renderer, factored out of m4b_gpu.py
so the offline (m4b_gpu) and live (live_render) paths share ONE render.

Faithful to the M5 CPU reference bit-for-bit (proven in M4B-RESULTS.md):
per-draw primitives -> GPU compute raster (6-edge clip, fp64 planes, nearest-z
winner) -> the verified perspective divide + real-texture texel decode as one
vectorized numpy post-pass.

  Renderer(ctx, texlist).frame(dump_range) -> HxWx3 uint8
where dump_range(lo, hi) -> bytes is emu860c.dump_range (a bulk memory read).
where r32(addr)->uint32 reads the running firmware's memory (the live per-draw
coefficient program lives at 0x08158000..0x08170000).
"""
import struct
import numpy as np
from dpl_sampler import mode_to_texflags, wrap_index, composite_masked

W, H = 832, 512
PROG_LO, PROG_HI = 0x08158000, 0x08170000
EDGE = {0x42, 0x0d, 0x2c}

FX_SHADER = """
#version 430
#extension GL_ARB_gpu_shader_fp64 : enable
layout(local_size_x = 64, local_size_y = 1) in;
layout(std430, binding = 0) readonly buffer Prims { double prims[]; };
// 11 dvec4 / prim (44 doubles): [0..5] edges (up to 6, A B C used),
//   [6] zdepth+nedges, [7] u+hasuv, [8] v+tid, [9] tz+0,
//   [10] litR,litG,litB,0 -- the polygon's TREEclmpintoMEM lit color
//   (dvpx_r24/g24/b24; DIVPXMAP.H), CONSTANT per polygon (not a plane).
layout(std430, binding = 1) buffer Outs { uint outbuf[]; };  // 9 words/pixel
uniform int n_prims;
uniform int width;
void main() {
    int x = int(gl_GlobalInvocationID.x);
    int y = int(gl_GlobalInvocationID.y);
    if (x >= width) return;
    double fx = double(x), fy = double(y);
    double zbest = -1e30lf;
    double uu = 0.0lf, vv = 0.0lf, tz = 1.0lf;
    float tid = 0.0, hasuv = 0.0;
    float litR = 255.0, litG = 255.0, litB = 255.0;
    for (int p = 0; p < n_prims; p++) {
        int b = p*44;
        bool allpos = true, allneg = true;
        for (int e = 0; e < 6; e++) {
            int eb = b + e*4;
            if (prims[eb+3] < 0.5lf) continue;
            double ve = prims[eb]*fx + prims[eb+1]*fy + prims[eb+2];
            if (ve >= 0.0lf) allneg = false; else allpos = false;
        }
        if (!(allpos || allneg)) continue;
        double z = prims[b+24]*fx + prims[b+25]*fy + prims[b+26];
        if (z <= zbest) continue;
        zbest = z;
        hasuv = float(prims[b+31]);
        uu = prims[b+28]*fx + prims[b+29]*fy + prims[b+30];
        vv = prims[b+32]*fx + prims[b+33]*fy + prims[b+34];
        tz = prims[b+36]*fx + prims[b+37]*fy + prims[b+38];
        tid = float(prims[b+35]);
        litR = float(prims[b+40]); litG = float(prims[b+41]); litB = float(prims[b+42]);
    }
    int pix = y*width + x;
    outbuf[pix*9+0] = (zbest > -1e29lf) ? 0x3f800000u : 0xff800000u;
    outbuf[pix*9+1] = floatBitsToUint(float(uu));
    outbuf[pix*9+2] = floatBitsToUint(float(vv));
    outbuf[pix*9+3] = floatBitsToUint(float(tz));
    outbuf[pix*9+4] = floatBitsToUint(tid);
    outbuf[pix*9+5] = floatBitsToUint(hasuv);
    outbuf[pix*9+6] = floatBitsToUint(litR);
    outbuf[pix*9+7] = floatBitsToUint(litG);
    outbuf[pix*9+8] = floatBitsToUint(litB);
}
"""


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


def build_prims(dump_range):
    """Reconstruct per-draw primitive records from the live program in memory.

    dump_range(lo, hi) -> bytes: ONE bulk read of the whole program window,
    replacing a ~24k-call Python r32() loop (the dominant per-frame cost that
    made the live renderer fall behind the running game -- see M4B-RESULTS.md).
    """
    words = np.frombuffer(dump_range(PROG_LO, PROG_HI), dtype='<u4')
    n = len(words)

    def rd(a):
        idx = (a - PROG_LO) >> 2
        return int(words[idx]) if 0 <= idx < n else 0

    setups = (PROG_LO + 4 * np.flatnonzero(words == 0x100)).tolist()
    recs4 = []
    for a in setups:
        p = a + 4
        edges = []
        while ((rd(p) >> 8) & 0xff) in EDGE and len(edges) < 6:
            edges.append((f32(rd(p + 4)), f32(rd(p + 8)), f32(rd(p + 12)))); p += 16
        if len(edges) < 3:
            continue
        zp = tzp = up = vp = None; tid = None; q = p
        # r24/g24/b24 (DIVPXMAP.H): the polygon's lit color, written via
        # TREEclmpintoMEM (op 0x5a) as a CONSTANT (3rd word of the 3-word
        # instruction), not a linear plane -- confirmed live: addr 118/126/134
        # (DIVPXMAP.H names them 117/125/133; the +1 is this bytecode's own
        # addressing, taken empirically from the wire, not the header comment).
        litr = litg = litb = None
        for _ in range(70):
            w = rd(q); op = (w >> 8) & 0xff; ad = w & 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 == 32 and tzp is None:
                tzp = (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 op == 0xf7 and ad in (141, 142) and tid is None:
                tid = (rd(q + 4) >> 2) & 0x3f
            if op == 0x5a and ad == 118 and litr is None:
                litr = f32(rd(q + 8))
            if op == 0x5a and ad == 126 and litg is None:
                litg = f32(rd(q + 8))
            if op == 0x5a and ad == 134 and litb is None:
                litb = f32(rd(q + 8))
            q += 4
        if zp is None:
            continue
        hasuv = 1.0 if (up and vp and tzp) else 0.0
        up = up or (0, 0, 0); vp = vp or (0, 0, 0); tzp = tzp or (0, 0, 1)
        # clamp direct to [0,255] -- the observed range (~49..395) is already
        # in 8-bit brightness units, not a [0,1] float to rescale (see
        # M4B-RESULTS.md "lit color plane" note); default 255 = no-op tint
        # for polys where the color plane wasn't found (unchanged behavior).
        litr = 255.0 if litr is None else max(0.0, min(255.0, litr))
        litg = 255.0 if litg is None else max(0.0, min(255.0, litg))
        litb = 255.0 if litb is None else max(0.0, min(255.0, litb))
        eslots = [(e[0], e[1], e[2], 1.0) for e in edges[:6]]
        while len(eslots) < 6:
            eslots.append((0.0, 0.0, 0.0, 0.0))
        recs4 += eslots + [(*zp, float(len(edges))), (*up, hasuv),
                           (*vp, float(tid or 0)), (*tzp, 0.0),
                           (litr, litg, litb, 0.0)]
    if not recs4:
        return None
    return np.array(recs4, dtype=np.float64).reshape(-1, 4)


class Renderer:
    def __init__(self, ctx, texlist):
        self.ctx = ctx
        self.prog_gl = ctx.compute_shader(FX_SHADER)
        self.texlist = texlist

    def frame(self, dump_range):
        prims = build_prims(dump_range)
        if prims is None:
            return None
        n_prims = len(prims) // 11
        b0 = self.ctx.buffer(prims.tobytes())
        out = self.ctx.buffer(reserve=W * H * 9 * 4)
        b0.bind_to_storage_buffer(0)
        out.bind_to_storage_buffer(1)
        self.prog_gl['n_prims'] = n_prims
        self.prog_gl['width'] = W
        self.prog_gl.run(group_x=(W + 63) // 64, group_y=H)
        raw = np.frombuffer(out.read(), dtype=np.uint32).reshape(H, W, 9)
        b0.release(); out.release()
        zb = raw[..., 0].view(np.float32)
        uu = raw[..., 1].view(np.float32).astype(np.float64)
        vv = raw[..., 2].view(np.float32).astype(np.float64)
        tz = raw[..., 3].view(np.float32).astype(np.float64)
        tid = raw[..., 4].view(np.float32).astype(np.int64)
        hasuv = raw[..., 5].view(np.float32)
        # lit color (dvpx_r24/g24/b24 -> TREEclmpintoMEM): the polygon's
        # computed lighting, clamped to [0,255] -- see build_prims. Used as
        # the direct color for flat polys, and as a multiplicative tint on
        # sampled texture color for textured ones.
        lit = np.stack([raw[..., 6].view(np.float32),
                        raw[..., 7].view(np.float32),
                        raw[..., 8].view(np.float32)], axis=-1)
        filled = zb > -1e29
        img = np.zeros((H, W, 3), np.uint8)
        flat = filled & (hasuv <= 0.5)
        img[flat] = np.clip(lit[flat], 0, 255).astype(np.uint8)
        tzc = np.where(np.abs(tz) < 1.0, np.sign(tz) + (tz == 0), tz)
        ucoord = (uu / tzc * 2048).astype(np.int64)
        vcoord = (vv / tzc * 2048).astype(np.int64)
        ntex = len(self.texlist)
        textured = filled & (hasuv > 0.5)
        tidm = tid % ntex                     # hoisted: was recomputed per slot
        for slot in range(ntex):
            sel = textured & (tidm == slot)
            if not sel.any():
                continue
            h, (tu, tv, mode, arr) = self.texlist[slot]
            wrap_u, wrap_v, cut_mode = mode_to_texflags(mode)
            # sample only the pixels THIS texture covers, not the whole frame
            # (composite_masked -- verified equivalent to the full-frame form
            # in dpl_sampler's conformance test; this loop runs once per
            # texture, so full-frame sampling here was O(ntex x W x H))
            ui = wrap_index((ucoord[sel] * tu) >> 8, tu, wrap_u)
            vi = wrap_index((vcoord[sel] * tv) >> 8, tv, wrap_v)
            samp = arr[vi, ui].astype(np.float64) * (lit[sel] / 255.0)
            composite_masked(img, sel, np.clip(samp, 0, 255).astype(np.uint8), cut_mode)
        return img