#!/usr/bin/env python3
"""retroDE_ps2 — Ch349 step 4: software reference image for an ACTUAL SH3 draw.

Rasterizes the chosen draw's REAL geometry (its TRI_STRIP, real per-vertex screen XY + perspective S/T/Q +
vertex RGBA) sampling the texture RECONSTRUCTED from GS local memory (gs_sh3_recon / gs_localmem). Perspective-
correct: S,T,Q (= s/w, t/w, 1/w) interpolate linearly in screen space, then u=(S/Q)*TW, v=(T/Q)*TH; texel ->
CLUT -> ABGR; TFX=MODULATE applies the vertex color (×2/255-ish, GS 128=1.0). This is the host-first reference
that must pixel-check against the real PCSX2 frame BEFORE anything goes to feeder/board.

Usage:  gs_sh3_draw_ref.py <dump.gs.zst> [--draw-idx N] [--tbp T --cbp C --tbw W --tw 512 --th 512]
                                          [--clut-order grid|linear] [--modulate 0|1] [--out DIR]
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gs_sh3_draw_census as C
import gs_sh3_recon as RC

SW, SH = 640, 480

def edge(ax, ay, bx, by, px, py):
    return (px-ax)*(by-ay) - (py-ay)*(bx-ax)

def raster_strip(verts, prim_type, tex_idx, pal, tw, th, modulate):
    """Rasterize a TRI_STRIP/TRI_FAN/TRIANGLE list of verts into an SW*SH RGBA buffer with a z-buffer."""
    fb = [(0,0,0,0)]*(SW*SH)
    zb = [-1]*(SW*SH)
    def tri(i0, i1, i2):
        v0, v1, v2 = verts[i0], verts[i1], verts[i2]
        x0,y0 = v0["x"],v0["y"]; x1,y1 = v1["x"],v1["y"]; x2,y2 = v2["x"],v2["y"]
        minx = max(0, int(min(x0,x1,x2))); maxx = min(SW-1, int(max(x0,x1,x2))+1)
        miny = max(0, int(min(y0,y1,y2))); maxy = min(SH-1, int(max(y0,y1,y2))+1)
        area = edge(x0,y0, x1,y1, x2,y2)
        if abs(area) < 1e-6: return
        inv = 1.0/area
        for py in range(miny, maxy+1):
            for px in range(minx, maxx+1):
                cx, cy = px+0.5, py+0.5
                w0 = edge(x1,y1, x2,y2, cx,cy)
                w1 = edge(x2,y2, x0,y0, cx,cy)
                w2 = edge(x0,y0, x1,y1, cx,cy)
                # inside test (either winding)
                if not ((w0>=0 and w1>=0 and w2>=0) or (w0<=0 and w1<=0 and w2<=0)):
                    continue
                b0, b1, b2 = w0*inv, w1*inv, w2*inv
                S = b0*v0["s"] + b1*v1["s"] + b2*v2["s"]
                T = b0*v0["t"] + b1*v1["t"] + b2*v2["t"]
                Q = b0*v0["q"] + b1*v1["q"] + b2*v2["q"]
                if abs(Q) < 1e-12: continue
                u = (S/Q)*tw; v = (T/Q)*th
                tx = int(u) % tw; ty = int(v) % th          # REPEAT wrap (SH3 CLAMP default)
                if tx < 0: tx += tw
                if ty < 0: ty += th
                pidx = tex_idx[ty*tw + tx]
                p = pal[pidx & 0xFF]
                r,g,b = p&0xFF, (p>>8)&0xFF, (p>>16)&0xFF
                if modulate:
                    vr = b0*((v0["rgba"])&0xFF)+b1*((v1["rgba"])&0xFF)+b2*((v2["rgba"])&0xFF)
                    vg = b0*((v0["rgba"]>>8)&0xFF)+b1*((v1["rgba"]>>8)&0xFF)+b2*((v2["rgba"]>>8)&0xFF)
                    vb = b0*((v0["rgba"]>>16)&0xFF)+b1*((v1["rgba"]>>16)&0xFF)+b2*((v2["rgba"]>>16)&0xFF)
                    r = min(255, int(r*vr/128.0)); g = min(255, int(g*vg/128.0)); b = min(255, int(b*vb/128.0))
                z = int(b0*v0["z"] + b1*v1["z"] + b2*v2["z"])
                o = py*SW + px
                if z >= zb[o]:
                    zb[o] = z; fb[o] = (r,g,b,255)
    n = len(verts)
    if prim_type == 3:            # TRIANGLE list
        for i in range(0, n-2, 3): tri(i,i+1,i+2)
    elif prim_type == 4:          # TRI_STRIP
        for i in range(2, n): tri(i-2, i-1, i)
    elif prim_type == 5:          # TRI_FAN
        for i in range(2, n): tri(0, i-1, i)
    return fb

def main(argv):
    if len(argv) < 2:
        print(__doc__); return 2
    dump = argv[1]
    def opt(n, dv=None): return argv[argv.index(n)+1] if n in argv else dv
    draw_idx = int(opt("--draw-idx","89761"))
    tbp = int(opt("--tbp","9216")); cbp = int(opt("--cbp","13952"))
    fbw = int(opt("--tbw","8")); tw = int(opt("--tw","512")); th = int(opt("--th","512"))
    order = opt("--clut-order","grid"); modulate = int(opt("--modulate","1"))
    outdir = opt("--out", os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                                       "captures","gs","silenthill3","extracted","recon"))
    os.makedirs(outdir, exist_ok=True)

    dr = C.get_draw(dump, draw_idx)
    if dr is None:
        print(f"draw first_idx={draw_idx} not found"); return 1
    t0 = dr["tex0"]
    print(f"[step4] draw f{dr['frame']} idx{draw_idx} {C.PRIMT[dr['prim']['type']]} nprim={dr['nprim']} "
          f"nvert={dr['nvert']} TEX0 tbp={t0['tbp']} psm=0x{t0['psm']:02x} {t0['tw']}x{t0['th']}")

    mem, replayed, uploads, events, vram = RC.build_localmem_to(dump, draw_idx)
    if mem is None: print("VRAM absent"); return 1
    tex_idx = mem.read_psmt8(tbp, fbw, tw, th)
    pal = RC.read_clut32(mem, cbp, order=order)
    print(f"[step4] reconstructed tex {tw}x{th} ({len(set(tex_idx))} idx), CLUT {order} ({len(set(pal))} colors)")

    fb = raster_strip(dr["verts"], dr["prim"]["type"], tex_idx, pal, tw, th, modulate)
    painted = sum(1 for p in fb if p[3])
    RC.save_png(os.path.join(outdir, f"draw_ref_{draw_idx}_{order}.png"), SW, SH, fb)
    print(f"[step4] painted {painted} px -> draw_ref_{draw_idx}_{order}.png  ({outdir})")
    # crop to the draw's on-screen bbox for easier visual compare with the PCSX2 frame
    x0=max(0,int(dr['xmin'])); x1=min(SW,int(dr['xmax'])+1); y0=max(0,int(dr['ymin'])); y1=min(SH,int(dr['ymax'])+1)
    if x1>x0 and y1>y0:
        crop = [fb[y*SW+x] for y in range(y0,y1) for x in range(x0,x1)]
        RC.save_png(os.path.join(outdir, f"draw_ref_{draw_idx}_{order}_crop.png"), x1-x0, y1-y0, crop)
        print(f"[step4] bbox crop x[{x0}..{x1}] y[{y0}..{y1}] -> draw_ref_{draw_idx}_{order}_crop.png")
    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv))