#!/usr/bin/env python3
"""retroDE_ps2 — Ch349 step 1: census of ACTUAL drawn textured geometry in a GS dump.

Ch347/348 proved authentic ASSETS (SH3 PSMT8 tex + real CLUT) through CHOSEN geometry. The remaining
authenticity gap (Codex Ch349): reconstruct an ACTUAL commercial draw faithfully — the texture as the real
draw samples it (streamed in one format, sampled in another) on the real triangle's ST/Q + screen geometry.

This tool is step 1: walk the GS register stream, reconstruct every textured drawing primitive with its full
per-vertex state (screen XY from XYZF2/XYZ2 12.4 fixed, S/T/Q from ST+RGBAQ.Q, RGBA), group consecutive
primitives that share TEX0+PRIM state into DRAWS, and rank them so a single real environment draw can be
PICKED for reconstruction. Pure stdlib; reuses gs_parse for the GIF/register walk and gs_texture_residency
for the VRAM snapshot + CLUT/texture residency verdict.

A good Ch349 candidate is: TME=1, texture RESIDENT in the VRAM snapshot, a non-trivial on-screen footprint
(a real surface, not a 2px HUD glyph), indexed or CT texture with a known PSM, and enough triangles to be a
genuine mapped surface. The census REPORTS; it does not pick for you — the ranked head is the shortlist.

Usage:  gs_sh3_draw_census.py <dump.gs.zst> [--top N] [--frame F] [--json out.json] [--min-prims K]
"""
import sys, os, json, struct
from collections import defaultdict
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gs_parse
import gs_texture_residency as R

def f32(bits):
    return struct.unpack("<f", struct.pack("<I", bits & 0xFFFFFFFF))[0]

PRIMT = {0:"POINT",1:"LINE",2:"LINE_STRIP",3:"TRIANGLE",4:"TRI_STRIP",5:"TRI_FAN",6:"SPRITE",7:"INVALID"}
VERTS_PER = {0:1,1:2,2:2,3:3,4:3,5:3,6:2,7:0}   # min verts to kick a primitive

def census(dump, frame_filter=None, min_prims=1, collected=None):
    d, h, events, uploads, runs, vram = collected if collected is not None else R.collect(dump, 0)
    # live GS state we latch as we walk
    prim = dict(type=7, tme=0, fst=0, abe=0, ctxt=0)
    tex0 = {1:None, 2:None}
    ofx = {1:0.0, 2:0.0}; ofy = {1:0.0, 2:0.0}
    cur_st = (0.0, 0.0, 1.0)   # S, T, Q
    cur_rgba = 0
    vqueue = []                # vertices latched since last primitive (for fan/strip we only need a window)
    draws = []
    cur = None

    def newkey():
        t0 = tex0[1 if prim["ctxt"]==0 else 2]
        if t0 is None: return None
        return (t0["tbp"], t0["psm"], t0["tbw"], prim["type"], prim["tme"], prim["abe"])

    def close():
        nonlocal cur
        if cur and cur["nprim"] >= 1:
            draws.append(cur)
        cur = None

    def open_draw():
        nonlocal cur
        t0 = tex0[1 if prim["ctxt"]==0 else 2]
        cur = dict(key=newkey(), tex0=dict(t0), prim=dict(prim), frame=cur_frame,
                   first_idx=cur_idx, nprim=0, nvert=0,
                   xmin=1e30, xmax=-1e30, ymin=1e30, ymax=-1e30,
                   smin=1e30, smax=-1e30, tmin=1e30, tmax=-1e30,
                   qmin=1e30, qmax=-1e30, verts=[])

    cur_frame = 0; cur_idx = 0
    for e in events:
        if e.kind == "FRAME_BOUNDARY":
            cur_frame = e.frame + 1
            continue
        if e.kind != "GSREG":
            continue
        cur_frame = e.frame; cur_idx = e.idx
        r, v = e.reg, e.value
        if r == "PRIM":
            close()
            prim = dict(type=v&7, tme=(v>>4)&1, fst=(v>>8)&1, abe=(v>>6)&1, ctxt=(v>>9)&1)
            vqueue = []
        elif r == "PRMODE":      # PRIM-less prim mode (rare); ignore topology change w/o reset
            pass
        elif r == "TEX0_1": tex0[1] = R.dec_tex0(v)
        elif r == "TEX0_2": tex0[2] = R.dec_tex0(v)
        elif r == "XYOFFSET_1": ofx[1] = (v & 0xFFFF)/16.0; ofy[1] = ((v>>32)&0xFFFF)/16.0
        elif r == "XYOFFSET_2": ofx[2] = (v & 0xFFFF)/16.0; ofy[2] = ((v>>32)&0xFFFF)/16.0
        elif r == "RGBAQ": cur_rgba = v & 0xFFFFFFFF
        elif r == "ST":
            s = f32(v & 0xFFFFFFFF); t = f32((v>>32)&0xFFFFFFFF)
            q = f32(e.info.get("q_stq", 0x3F800000))   # Q rides in ST lane2; default 1.0
            cur_st = (s, t, q)
        elif r in ("XYZF2","XYZ2","XYZF3","XYZ3"):
            # drawing kick. XY are 12.4 fixed relative to XYOFFSET.
            xf = (v & 0xFFFF); yf = (v>>16) & 0xFFFF
            ci = 1 if prim["ctxt"]==0 else 2
            x = xf/16.0 - ofx[ci]; y = yf/16.0 - ofy[ci]
            if r in ("XYZF2","XYZF3"):
                z = (v>>32) & 0xFFFFFF
            else:
                z = (v>>32) & 0xFFFFFFFF
            if not prim["tme"]:           # only textured draws are Ch349 candidates
                continue
            if newkey() is None:
                continue
            if cur is None or cur["key"] != newkey():
                close(); open_draw()
            vtx = dict(x=x, y=y, z=z, s=cur_st[0], t=cur_st[1], q=cur_st[2], rgba=cur_rgba)
            cur["verts"].append(vtx); cur["nvert"] += 1
            cur["nprim"] += 1   # each kick completes one primitive in fan/strip/list
            cur["xmin"]=min(cur["xmin"],x); cur["xmax"]=max(cur["xmax"],x)
            cur["ymin"]=min(cur["ymin"],y); cur["ymax"]=max(cur["ymax"],y)
            for (a,lo,hi) in (("s","smin","smax"),("t","tmin","tmax"),("q","qmin","qmax")):
                pass
            cur["smin"]=min(cur["smin"],cur_st[0]); cur["smax"]=max(cur["smax"],cur_st[0])
            cur["tmin"]=min(cur["tmin"],cur_st[1]); cur["tmax"]=max(cur["tmax"],cur_st[1])
            cur["qmin"]=min(cur["qmin"],cur_st[2]); cur["qmax"]=max(cur["qmax"],cur_st[2])
    close()

    # attach residency verdict + derived UV/on-screen metrics + score; filter
    SW, SH = 640.0, 480.0   # SH3 internal render res (matches header ss=640x480)
    out = []
    for dr in draws:
        if dr["nprim"] < min_prims: continue
        if frame_filter is not None and dr["frame"] != frame_filter: continue
        t0 = dr["tex0"]
        snap = R.snapshot_present(vram, t0["tbp"], nb=512)
        clut = None
        if t0["psm"] in R.INDEXED_PSMS:
            csnap = R.snapshot_present(vram, t0["cbp"], nb=1024, min_nz=64)
            clut = dict(cbp=t0["cbp"], cpsm=t0["cpsm"], cld=t0["cld"], resident=bool(csnap),
                        distinct=(csnap["distinct"] if csnap else None))
        dr["tex_resident"] = bool(snap)
        dr["tex_snap"] = snap
        dr["clut"] = clut
        dr["w_px"] = dr["xmax"]-dr["xmin"]; dr["h_px"] = dr["ymax"]-dr["ymin"]
        # per-vertex perspective UV (u_norm=S/Q) -> texel bbox; on-screen vertex fraction
        umin=vmin=1e30; umax=vmax=-1e30; onscreen=0
        for vtx in dr["verts"]:
            if 0.0 <= vtx["x"] <= SW and 0.0 <= vtx["y"] <= SH: onscreen += 1
            q = vtx["q"]
            if abs(q) < 1e-9: continue
            un = vtx["s"]/q; vn = vtx["t"]/q
            umin=min(umin,un); umax=max(umax,un); vmin=min(vmin,vn); vmax=max(vmax,vn)
        dr["onscreen_frac"] = onscreen/max(1,dr["nvert"])
        if umax >= umin:
            dr["u_texmin"]=umin*t0["tw"]; dr["u_texmax"]=umax*t0["tw"]
            dr["v_texmin"]=vmin*t0["th"]; dr["v_texmax"]=vmax*t0["th"]
        else:
            dr["u_texmin"]=dr["u_texmax"]=dr["v_texmin"]=dr["v_texmax"]=0.0
        # on-screen clipped bbox area
        cx0=max(0.0,dr["xmin"]); cx1=min(SW,dr["xmax"]); cy0=max(0.0,dr["ymin"]); cy1=min(SH,dr["ymax"])
        dr["onscreen_area"]=max(0.0,cx1-cx0)*max(0.0,cy1-cy0)
        dr["score"] = _score(dr)
        out.append(dr)
    out.sort(key=lambda x: x["score"], reverse=True)
    return out, h, vram

def _score(dr):
    """'Good Ch349 candidate' = resident textured surface, MOSTLY ON-SCREEN, sampling a real texel
    rectangle in perspective. Reward on-screen containment + sampled-texel span, NOT guard-band area."""
    t0 = dr["tex0"]
    s = 0.0
    if not dr["tex_resident"]: return -1.0          # must be reconstructable
    if t0["psm"] in R.INDEXED_PSMS:
        if dr["clut"] and dr["clut"]["resident"]: s += 40.0   # indexed + resident CLUT == the real SH3 path
        else: return -1.0
    elif t0["psm"] in (0x00,0x02,0x0A,0x01):        # CT32/CT16/CT16S/CT24 — directly decodable
        s += 18.0
    else:
        return -1.0                                  # unsupported PSM for host decode
    # ON-SCREEN containment is the dominant term (we want a draw we can actually show + check)
    s += 100.0 * dr["onscreen_frac"]
    s += min(dr["onscreen_area"]/200.0, 120.0)       # on-screen area only (guard band excluded)
    # sampled texel rectangle must be a real chunk, not a single degenerate texel
    du = abs(dr["u_texmax"]-dr["u_texmin"]); dv = abs(dr["v_texmax"]-dr["v_texmin"])
    s += min(du, 64.0) + min(dv, 64.0)
    if du < 1.0 and dv < 1.0: s -= 80.0              # near-constant UV: flat/degenerate, not interesting
    s += min(dr["nprim"], 48)
    if dr["qmax"] > dr["qmin"] * 1.02: s += 15.0     # genuine perspective
    return s

def fmt(dr):
    t0 = dr["tex0"]; p = dr["prim"]
    clut = ""
    if dr["clut"]:
        clut = f" CLUT[{'R' if dr['clut']['resident'] else 'X'} cbp={t0['cbp']} cpsm=0x{t0['cpsm']:02x} cld={t0['cld']} dist={dr['clut']['distinct']}]"
    persp = "PERSP" if dr["qmax"] > dr["qmin"]*1.02 else "affine"
    return (f"score={dr['score']:6.1f} f{dr['frame']} idx{dr['first_idx']} {PRIMT[p['type']]} "
            f"tme={p['tme']} abe={p['abe']} fst={p['fst']} nprim={dr['nprim']} onscr={dr['onscreen_frac']*100:.0f}%\n"
            f"    TEX0 tbp={t0['tbp']} tbw={t0['tbw']} psm=0x{t0['psm']:02x} {t0['tw']}x{t0['th']} "
            f"tcc={t0['tcc']} tfx={t0['tfx']} resident={dr['tex_resident']}{clut}\n"
            f"    screen x[{dr['xmin']:.1f}..{dr['xmax']:.1f}] y[{dr['ymin']:.1f}..{dr['ymax']:.1f}] "
            f"on-area={dr['onscreen_area']:.0f}px2  texel u[{dr['u_texmin']:.1f}..{dr['u_texmax']:.1f}] "
            f"v[{dr['v_texmin']:.1f}..{dr['v_texmax']:.1f}] {persp}")

def get_draw(dump, first_idx):
    """Return the single census draw whose first_idx matches (with its full vertex list), or None."""
    draws, h, vram = census(dump, frame_filter=None, min_prims=1)
    for dr in draws:
        if dr["first_idx"] == first_idx:
            return dr
    return None

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
    top = int(opt("--top","25")); minp = int(opt("--min-prims","1"))
    ff = int(opt("--frame")) if "--frame" in argv else None
    draws, h, vram = census(dump, frame_filter=ff, min_prims=minp)
    print(f"# Ch349 draw census: {os.path.basename(dump)}")
    print(f"# textured draws (>= {minp} prim): {len(draws)}  vram_snapshot={'present' if vram is not None else 'ABSENT'}")
    cand = [d for d in draws if d["score"] > 0]
    print(f"# reconstructable candidates (resident tex + known PSM): {len(cand)}\n")
    for dr in draws[:top]:
        print(fmt(dr)); print()
    if "--json" in argv:
        slim = [dict(score=d["score"], frame=d["frame"], first_idx=d["first_idx"],
                     prim=d["prim"], tex0=d["tex0"], nprim=d["nprim"],
                     screen=dict(xmin=d["xmin"],xmax=d["xmax"],ymin=d["ymin"],ymax=d["ymax"]),
                     st=dict(smin=d["smin"],smax=d["smax"],tmin=d["tmin"],tmax=d["tmax"],
                             qmin=d["qmin"],qmax=d["qmax"]),
                     tex_resident=d["tex_resident"], clut=d["clut"]) for d in draws]
        open(opt("--json"),"w").write(json.dumps(slim, indent=1)+"\n")
        print(f"# wrote {opt('--json')}")
    return 0

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