#!/usr/bin/env python3
"""retroDE_ps2 — Ch346 generic texture-residency preflight.

Stops the hand-chasing of stale-VRAM frames (Ch345b: 32 sprites bound TBP=13440, but the captured frame's
VRAM there held cube-checker residue, not the font). Before ANY repack/render/fit, this proves — for each
textured draw run in a dump — that the bound TEX0 maps to a REAL upload whose payload actually covers the
sampled footprint and looks like resident content, NOT stale/placeholder VRAM.

Generic (per Codex): the gate is RESIDENT + PLAUSIBLE content, not "font-like" — future targets may be
sprites, UI, backgrounds, or indexed textures. `--font-like` adds the glyph-specific extra check on top.

Checks (Codex's minimum set):
  1. active draw TEX0 maps to a known upload region (DBP/DPSM/stride), not just an address;
  2. upload EPOCH tracked — if a later upload overwrites that TBP, the candidate uses the latest payload;
  3. sampled footprint (UV bbox, REPEAT-wrapped into the TEX0 TW/TH) is COVERED by the uploaded rect;
  4. payload sanity — reject all-zero / single-color flat / flat-alpha-on-alpha-draw / known stale hashes;
  5. emit RANKED candidates with frame/event offsets, prim run, TEX0, upload source offset, PSM, dims,
     alpha stats, and WHY it passed/failed.
Repack/render tools refuse to run unless `residency_ok()` returns a PASS for the bound texture.

A checker is NOT auto-rejected — it can be legit authentic content (the Ch343 cube). The signal that killed
Ch345b is RESIDENCY (no upload to TBP=13440 at all), not "it's a checker".

Usage:  gs_texture_residency.py <dump.gs[.xz|.zst]> [--max-runs N] [--font-like] [--report r.txt] [--json j.json]
        gs_texture_residency.py <dump> --assert TBP[:PSM]   # exit 0 iff that TBP is resident+plausible
"""
import sys, os, json, hashlib
from collections import Counter
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gs_parse

PSMCT32 = 0x00
MIN_DISTINCT_COLORS = 2          # <2 => flat fill (stale/cleared); fonts (B/W+transp=3) still pass
KNOWN_STALE_HASHES = set()       # sha256[:16] denylist of textures known to be placeholder/test residue

INDEXED_PSMS = (0x13, 0x14)   # PSMT8, PSMT4 — need a resident CLUT palette to render
def dec_tex0(v):
    return dict(tbp=v&0x3FFF, tbw=(v>>14)&0x3F, psm=(v>>20)&0x3F,
                tw=1<<((v>>26)&0xF), th=1<<((v>>30)&0xF), tcc=(v>>34)&1, tfx=(v>>35)&3,
                cbp=(v>>37)&0x3FFF, cpsm=(v>>51)&0xF, csm=(v>>55)&1,
                csa=(v>>56)&0x1F, cld=(v>>61)&7)
def dec_bitbltbuf(v):
    return dict(sbp=v&0x3FFF, sbw=(v>>16)&0x3F, spsm=(v>>24)&0x3F,
                dbp=(v>>32)&0x3FFF, dbw=(v>>48)&0x3F, dpsm=(v>>56)&0x3F)
def dec_trxreg(v):  return ((v & 0xFFF), ((v>>32) & 0xFFF))          # RRW, RRH
def dec_trxpos(v):  return (((v>>32)&0x7FF), ((v>>48)&0x7FF))        # DSAX, DSAY
def dec_trxdir(v):  return v & 3                                     # 0 = host->local (upload)

def payload_stats(blob, psm):
    """Decode an upload payload and summarize. Only PSMCT32 fully supported; others -> stats=None."""
    if psm != PSMCT32 or len(blob) < 4:
        return None
    n = len(blob)//4
    texels = [int.from_bytes(blob[i*4:i*4+4], "little") for i in range(n)]
    colors = Counter(texels)
    a = [(t>>24)&0xFF for t in texels]
    transp = sum(1 for x in a if x == 0); opaque = sum(1 for x in a if x >= 0x80)
    partial = n - transp - opaque
    h = hashlib.sha256(blob).hexdigest()[:16]
    return dict(texels=n, distinct=len(colors), top=colors.most_common(1)[0][1],
                a_transp=transp, a_opaque=opaque, a_partial=partial,
                flat_alpha=(len(set(a))==1), sha16=h,
                checkerish=_checkerish(texels))

def _checkerish(texels):
    """Heuristic structural flag (REPORTED, never an auto-reject — a checker can be legit content):
    a 2-color image whose value flips on a regular coarse grid. Returns block size or 0."""
    vals = set(texels)
    if len(vals) > 4: return 0
    n = len(texels); side = int(n**0.5)
    if side*side != n: return 0
    rows = [texels[r*side:(r+1)*side] for r in range(side)]
    # find the run length of the first row; a checker has long uniform runs that alternate
    first = rows[0]; run = 1
    while run < side and first[run] == first[0]: run += 1
    if run < 4 or run > side//2: return 0
    return run

def collect(path, max_runs):
    d = gs_parse.read_dump_bytes(path)
    h, events = gs_parse.parse_dump(path)
    uploads = []                 # ordered upload log
    epoch = {}                   # dbp -> count seen so far
    bbuf = trxreg = trxpos = None; trxdir = 3
    for e in events:
        if e.kind == "GSREG":
            if   e.reg == "BITBLTBUF": bbuf = dec_bitbltbuf(e.value)
            elif e.reg == "TRXREG":    trxreg = dec_trxreg(e.value)
            elif e.reg == "TRXPOS":    trxpos = dec_trxpos(e.value)
            elif e.reg == "TRXDIR":    trxdir = dec_trxdir(e.value)
        elif e.kind == "IMAGE" and bbuf is not None and trxdir == 0:    # host->local upload
            w, hh = trxreg if trxreg else (0,0); dx, dy = trxpos if trxpos else (0,0)
            dbp = bbuf["dbp"]; epoch[dbp] = epoch.get(dbp, 0) + 1
            nbytes = e.info.get("bytes", 0)
            uploads.append(dict(idx=e.idx, frame=e.frame, byte_off=e.byte_off, bytes=nbytes,
                                dbp=dbp, dbw=bbuf["dbw"], dpsm=bbuf["dpsm"], w=w, h=hh, dx=dx, dy=dy,
                                epoch=epoch[dbp], blob_range=(e.byte_off, e.byte_off+nbytes)))
    # draw runs: contiguous textured (TME) draws sharing one TEX0 base+psm; track sampled UV bbox
    runs = _draw_runs(events, max_runs)
    # VRAM snapshot: PCSX2 GS dumps freeze the full 4 MiB GS local memory at the END of the state blob
    # (register prefix first, then VRAM). Commercial games upload textures/CLUTs at scene-load — BEFORE the
    # dump — so they live here, not as in-stream upload events. This is the correct "resident" source.
    VRAM = 0x400000
    vstart = h.packet_start - 8192 - VRAM
    vram = d[vstart:vstart+VRAM] if 0 <= vstart and vstart + VRAM <= len(d) else None
    return d, h, events, uploads, runs, vram

def snapshot_present(vram, base, nb=512, min_nz=16):
    """Is there resident, non-flat content at this base pointer (256-byte units) in the VRAM snapshot?
    The snapshot is SWIZZLED, so this is a presence/plausibility check, not accurate content (de-swizzle
    is the translation step). Returns stats dict or None."""
    if vram is None: return None
    o = base * 256
    if o < 0 or o + nb > len(vram): return None
    chunk = vram[o:o+nb]
    nz = sum(1 for b in chunk if b); dist = len(set(chunk))
    return dict(nonzero=nz, distinct=dist) if (nz >= min_nz and dist >= 2) else None

def _draw_runs(events, max_runs):
    cur = {"type":None,"tme":0,"fst":0,"abe":0}; tex0 = {1:None,2:None}; ctxt = 0
    uvbuf = []; runs = []; run = None
    def close():
        nonlocal run
        if run and run["nprim"] > 0: runs.append(run)
        run = None
    for e in events:
        if e.kind != "GSREG": continue
        r, v = e.reg, e.value
        if r == "PRIM":
            cur = {"type":v&7,"tme":(v>>4)&1,"fst":(v>>8)&1,"abe":(v>>6)&1}; ctxt=(v>>9)&1; uvbuf=[]
        elif r == "TEX0_1": tex0[1]=dec_tex0(v)
        elif r == "TEX0_2": tex0[2]=dec_tex0(v)
        elif r == "UV": uvbuf.append((v&0x3FFF,(v>>14)&0x3FFF))
        elif r == "XYZ2":
            if not cur["tme"]: continue                       # only textured draws make residency runs
            t0 = tex0[1 if ctxt==0 else 2]
            if t0 is None: continue
            key = (t0["tbp"], t0["psm"], t0["tbw"])
            if run is None or run["key"] != key or run["last_idx_gap"] != e.idx-1:
                close()
                run = dict(key=key, tex0=t0, type=cur["type"], abe=cur["abe"], fst=cur["fst"],
                           first_idx=e.idx, frame=e.frame, nprim=0,
                           umin=1<<30,umax=-1,vmin=1<<30,vmax=-1, last_idx_gap=e.idx)
                if len(runs) >= max_runs and max_runs>0: break
            run["nprim"] += 1; run["last_idx"]=e.idx; run["last_idx_gap"]=e.idx
            for (u,vv) in uvbuf[-2:]:                          # sprite=2 verts; tri uses last 3 — bbox is fine
                tu=u>>4; tv=vv>>4
                run["umin"]=min(run["umin"],tu); run["umax"]=max(run["umax"],tu)
                run["vmin"]=min(run["vmin"],tv); run["vmax"]=max(run["vmax"],tv)
            uvbuf=[]
    close()
    return runs

def evaluate(run, uploads, d, vram=None):
    """Return verdict dict: resident (in-stream upload OR VRAM snapshot) + coverage + plausibility + CLUT."""
    t0 = run["tex0"]; tbp=t0["tbp"]; psm=t0["psm"]; tbw=t0["tbw"]
    # candidate uploads: same base + psm, occurring BEFORE the run's first draw; pick latest by idx
    cands = [u for u in uploads if u["dbp"]==tbp and u["dpsm"]==psm and u["idx"] < run["first_idx"]]
    reasons = []
    if not cands:
        # no in-stream upload — fall back to the VRAM snapshot (scene-load uploads live there)
        snap = snapshot_present(vram, tbp, nb=512)
        if snap is None:
            any_dbp = any(u["dbp"]==tbp for u in uploads)
            reasons.append("texture NOT resident: no in-stream upload to TBP and VRAM snapshot empty/absent"
                           if not any_dbp else "upload(s) to TBP exist but PSM mismatch / all after draw, and snapshot empty")
            return dict(verdict="REJECT", reasons=reasons, upload=None, coverage=None, stats=None, clut=None, tex_source=None)
        reasons.append(f"texture resident in VRAM SNAPSHOT @tbp={tbp} (nz={snap['nonzero']}/512 distinct={snap['distinct']}; swizzled — content via translation)")
        clut = _clut_residency(t0, uploads, run, vram, reasons)
        clut_ok = (psm not in INDEXED_PSMS) or (clut is not None and clut["resident"])
        return dict(verdict=("PASS" if clut_ok else "REJECT"), reasons=reasons or ["resident (snapshot)"],
                    upload=None, coverage=None, stats=None, clut=clut, tex_source="snapshot")
    up = max(cands, key=lambda u: u["idx"])
    # coverage: sampled footprint, REPEAT-wrapped into TW/TH, must fall inside the uploaded rect [dx..dx+w)x[dy..dy+h)
    tw, th = t0["tw"], t0["th"]
    foot_known = run["umax"] >= 0 and run["vmax"] >= 0   # UV captured (fst=1); fst=0 ST/Q not yet sampled
    if not foot_known:
        # honest: do NOT claim coverage we didn't verify. Verdict rests on residency + plausibility.
        reasons.append("footprint UNVERIFIED (ST/Q draw — UV not captured); coverage not asserted")
        inside = True; coverage = None
    else:
        wrap = (run["umax"]>=tw or run["vmax"]>=th or run["umin"]<0 or run["vmin"]<0)
        fmin_u = run["umin"] % tw if tw else run["umin"]; fmin_v = run["vmin"] % th if th else run["vmin"]
        fmax_u = (run["umax"] % tw if tw else run["umax"]); fmax_v=(run["vmax"] % th if th else run["vmax"])
        inside = (up["dx"] <= fmin_u and fmax_u < up["dx"]+up["w"] and
                  up["dy"] <= fmin_v and fmax_v < up["dy"]+up["h"]) if (up["w"] and up["h"]) else False
        coverage = 1.0 if inside else 0.0
        if wrap: reasons.append(f"footprint WRAPS texture ({run['umin']}..{run['umax']} x {run['vmin']}..{run['vmax']} vs {tw}x{th}); REPEAT declared")
        if not inside: reasons.append(f"sampled footprint NOT covered by upload rect ({up['dx']}..{up['dx']+up['w']} x {up['dy']}..{up['dy']+up['h']})")
    blob = d[up["blob_range"][0]:up["blob_range"][1]]
    stats = payload_stats(blob, psm)
    plausible = True
    if stats is None:
        reasons.append(f"payload plausibility UNSUPPORTED for PSM 0x{psm:02x} (fail-closed)"); plausible=False
    else:
        if in_known_stale(stats["sha16"]): reasons.append(f"payload is a KNOWN stale/test texture ({stats['sha16']})"); plausible=False
        if stats["distinct"] < MIN_DISTINCT_COLORS: reasons.append(f"payload flat ({stats['distinct']} color)"); plausible=False
        if run["abe"] and stats["flat_alpha"]: reasons.append("alpha draw but payload alpha is FLAT (no mask)"); plausible=False
        if stats["checkerish"]: reasons.append(f"payload is structurally checker-like (block~{stats['checkerish']}) — verify it's intended content")
    clut = _clut_residency(t0, uploads, run, vram, reasons)
    clut_ok = (psm not in INDEXED_PSMS) or (clut is not None and clut["resident"])
    verdict = "PASS" if (inside and plausible and clut_ok) else "REJECT"
    return dict(verdict=verdict, reasons=reasons or ["resident + plausible"], upload=up, coverage=coverage, stats=stats, clut=clut, tex_source="stream")

def _clut_residency(t0, uploads, run, vram, reasons):
    """Ch347 — indexed textures (PSMT8/PSMT4) need a resident CLUT at CBP (in-stream upload OR VRAM
    snapshot). The datapath proof is NOT authentic ingestion: the emitted TEX0's CBP/CPSM/CLD must select
    a CLUT that is actually loaded."""
    if t0["psm"] not in INDEXED_PSMS: return None
    cbp = t0["cbp"]
    ccands = [u for u in uploads if u["dbp"] == cbp and u["idx"] < run["first_idx"]]
    if ccands:
        cup = max(ccands, key=lambda u: u["idx"])
        reasons.append(f"CLUT resident (stream) @cbp={cbp} cpsm=0x{t0['cpsm']:02x} cld={t0['cld']} (upload idx{cup['idx']} f{cup['frame']})")
        clut = dict(resident=True, source="stream", upload_idx=cup["idx"], frame=cup["frame"], cbp=cbp, cpsm=t0["cpsm"], cld=t0["cld"], distinct=None)
    else:
        snap = snapshot_present(vram, cbp, nb=1024, min_nz=64)    # a real 256-entry palette is rich + non-flat
        if snap:
            reasons.append(f"CLUT resident (snapshot) @cbp={cbp} cpsm=0x{t0['cpsm']:02x} cld={t0['cld']} (nz={snap['nonzero']}/1024 distinct={snap['distinct']})")
            clut = dict(resident=True, source="snapshot", upload_idx=None, frame=None, cbp=cbp, cpsm=t0["cpsm"], cld=t0["cld"], distinct=snap["distinct"])
        else:
            reasons.append(f"CLUT NOT resident: no upload to CBP={cbp} and snapshot empty — indexed texture cannot render authentically")
            clut = dict(resident=False, source=None, upload_idx=None, frame=None, cbp=cbp, cpsm=t0["cpsm"], cld=t0["cld"], distinct=None)
    if t0["cld"] == 0:
        reasons.append("CLD=0: this TEX0 does not trigger a CLUT (re)load — palette would be whatever a prior load left")
    return clut

def in_known_stale(sha16): return sha16 in KNOWN_STALE_HASHES   # helper (denylist seeded empty by default)

def font_like(stats):
    """Glyph-specific EXTRA check (only when --font-like): mask-like alpha (real transparent + opaque
    regions), modest palette. NOT part of the generic gate."""
    if stats is None: return False, "no stats"
    if stats["a_transp"] == 0: return False, "no transparent texels (not a glyph mask)"
    if stats["a_opaque"] == 0: return False, "no opaque texels"
    frac_t = stats["a_transp"]/stats["texels"]
    if not (0.05 <= frac_t <= 0.95): return False, f"transparent fraction {frac_t:.2f} not mask-like"
    return True, f"mask-like (transp {frac_t:.2f})"

def residency_ok(path, tbp, psm=PSMCT32):
    """Programmatic gate for repack/render tools. True iff SOME textured run binding (tbp,psm) is PASS."""
    d, h, events, uploads, runs, vram = collect(path, 0)
    for run in runs:
        if run["tex0"]["tbp"]==tbp and run["tex0"]["psm"]==psm:
            if evaluate(run, uploads, d, vram)["verdict"]=="PASS": return True
    return False

def main(argv):
    if len(argv) < 2: print(__doc__); return 2
    path = argv[1]
    def has(f): return f in argv
    def opt(n,dv=None): return argv[argv.index(n)+1] if n in argv else dv
    d, h, events, uploads, runs, vram = collect(path, int(opt("--max-runs","0")))

    if has("--assert"):
        spec = opt("--assert"); tbp = int(spec.split(":")[0],0); psm=int(spec.split(":")[1],0) if ":" in spec else PSMCT32
        ok = residency_ok(path, tbp, psm)
        print(f"residency {'PASS' if ok else 'REJECT'} for TBP={tbp} PSM=0x{psm:02x}")
        return 0 if ok else 1

    R = [f"# Ch346 texture-residency preflight: {os.path.basename(path)}",
         f"uploads(host->local): {len(uploads)}   textured draw runs: {len(runs)}"]
    results = []
    for run in runs:
        ev = evaluate(run, uploads, d, vram); t0=run["tex0"]
        fl = font_like(ev["stats"]) if has("--font-like") else None
        verdict = ev["verdict"]
        if verdict=="PASS" and fl is not None and not fl[0]: verdict="REJECT"
        results.append(dict(run=run, ev=ev, font=fl, verdict=verdict))
    passes = [r for r in results if r["verdict"]=="PASS"]
    passes.sort(key=lambda r: (r["ev"]["stats"]["distinct"] if r["ev"]["stats"] else 0, r["run"]["nprim"]), reverse=True)

    R.append(f"\n== {len(passes)} PASS / {len(results)-len(passes)} REJECT ==")
    for i,r in enumerate(results):
        run=r["run"]; ev=r["ev"]; t0=run["tex0"]; up=ev["upload"]
        R.append(f"\n[{r['verdict']}] run f{run['frame']} idx{run['first_idx']}+{run['nprim']}prim  "
                 f"TEX0 tbp={t0['tbp']} tbw={t0['tbw']} psm=0x{t0['psm']:02x} {t0['tw']}x{t0['th']} abe={run['abe']} fst={run['fst']}")
        if run["umax"] >= 0: R.append(f"    footprint u[{run['umin']}..{run['umax']}] v[{run['vmin']}..{run['vmax']}]")
        else:                R.append(f"    footprint UNVERIFIED (fst={run['fst']} ST/Q draw)")
        if up: R.append(f"    upload  @byte0x{up['byte_off']:x} f{up['frame']} epoch{up['epoch']} dbp={up['dbp']} dpsm=0x{up['dpsm']:02x} {up['w']}x{up['h']} @({up['dx']},{up['dy']}) {up['bytes']}B")
        if ev["stats"]:
            s=ev["stats"]; R.append(f"    payload distinct={s['distinct']} alpha[t/o/p]={s['a_transp']}/{s['a_opaque']}/{s['a_partial']} flatA={s['flat_alpha']} checker~{s['checkerish']} sha={s['sha16']}")
        if ev.get("clut") is not None:
            c=ev["clut"]; R.append(f"    CLUT {'RESIDENT' if c['resident'] else 'MISSING'} cbp={c['cbp']} cpsm=0x{c['cpsm']:02x} cld={c['cld']}")
        if r["font"] is not None: R.append(f"    font-like: {r['font'][0]} ({r['font'][1]})")
        for why in ev["reasons"]: R.append(f"    -> {why}")
    if passes:
        R.append(f"\n== RANKED PASS candidates ==")
        for r in passes:
            run=r["run"]; t0=run["tex0"]
            s=r['ev']['stats']; c=r['ev'].get('clut')
            extra = f"distinct={s['distinct']}" if s else f"tex_src={r['ev'].get('tex_source')}"
            if c: extra += f" CLUT@cbp={c['cbp']}/{c['source']}"
            R.append(f"  tbp={t0['tbp']} psm=0x{t0['psm']:02x} {run['nprim']}prim {extra} f{run['frame']} idx{run['first_idx']}")
    report = "\n".join(R)
    print(report)
    if has("--report"): open(opt("--report"),"w").write(report+"\n")
    if has("--json"):
        j = [dict(verdict=r["verdict"], tbp=r["run"]["tex0"]["tbp"], psm=r["run"]["tex0"]["psm"],
                  frame=r["run"]["frame"], first_idx=r["run"]["first_idx"], nprim=r["run"]["nprim"],
                  stats=r["ev"]["stats"], clut=r["ev"].get("clut"), reasons=r["ev"]["reasons"]) for r in results]
        open(opt("--json"),"w").write(json.dumps(j, indent=1)+"\n")
    return 0 if passes else 1

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