#!/usr/bin/env python3
"""retroDE_ps2 — Ch344 authentic SPRITE extractor (brick 1).

Selects the EARLIEST contiguous run of v1-eligible sprites from a .gs dump and emits a structured sprite
list + aggregate report. v1 eligibility (per the Ch344 census-gated scope) — FAIL CLOSED on anything else:
  * SPRITE primitive (PRIM type 6), TME=1 (textured), FST=1 (UV/affine)
  * TEX0 PSM = PSMCT32 (no CLUT / no PSMT8/PSMT4 / no PSMCT16)
  * ABE=1 (alpha) — blend equation is DECLARED source-over (ALPHA lives in the GS freeze state, absent
    from the packet stream; we emit ALPHA ourselves to the feeder)
  * small: width,height <= MAX_SPRITE_PX (excludes the fullscreen/scissored/guard-band blits)
  * single shared TEX0 TBP across the run (a TBP change / re-upload ends the run — never silently mixed)

Output is the sprite geometry/UV/color (dump-derived -> the .sprites file is LOCAL) plus an AGGREGATE
report (committable). This is extraction ONLY; rendering/translation to the feeder is brick 2+.

Usage:  gs_extract_sprites.py <dump.gs[.xz|.zst]> [--max N] [--out scene.sprites] [--report r.txt]
"""
import sys, os
from collections import Counter
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from gs_parse import parse_dump

PSMCT32 = 0x00
MAX_SPRITE_PX = 64
DEF_MAX = 32

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,
                cpsm=(v>>51)&0xF, csm=(v>>55)&1)

def extract(events, max_spr):
    cur = {"type":None,"tme":0,"abe":0,"fst":0}
    tex0 = {1:None, 2:None}; xyoff = {1:(0,0), 2:(0,0)}; ctxt = 0
    uvbuf = []; xyzbuf = []
    seg = []; seg_tbp = None; started = False; stop = None; first = last = -1
    rejected = Counter()

    def eligible(t0):
        return t0 is not None and cur["tme"] and cur["fst"] and t0["psm"]==PSMCT32 and cur["abe"] \
               and t0["csm"]==0 and t0["cpsm"]==0   # csm/cpsm: no CLUT in play

    for e in events:
        if e.kind != "GSREG": continue
        a, v = e.addr, e.value
        if a == 0x00:
            cur = {"type":v&7,"tme":(v>>4)&1,"abe":(v>>6)&1,"fst":(v>>8)&1}; ctxt=(v>>9)&1
            uvbuf=[]; xyzbuf=[]
        elif a in (0x06,0x07): tex0[1 if a==0x06 else 2] = dec_tex0(v)
        elif a in (0x18,0x19): xyoff[1 if a==0x18 else 2] = (v&0xFFFF,(v>>32)&0xFFFF)
        elif a == 0x03:                                   # UV
            uvbuf.append((v & 0x3FFF, (v>>14) & 0x3FFF))  # U,V in 10.4 fixed (14b)
        elif a == 0x05:                                   # XYZ2 vertex kick
            if cur["type"] != 6: continue
            xyzbuf.append(v)
            if len(xyzbuf) < 2: continue
            cx = 1 if ctxt==0 else 2; t0 = tex0[cx]; ox, oy = xyoff[cx]
            x0=((xyzbuf[0]&0xFFFF)-ox)//16; y0=(((xyzbuf[0]>>16)&0xFFFF)-oy)//16
            x1=((xyzbuf[1]&0xFFFF)-ox)//16; y1=(((xyzbuf[1]>>16)&0xFFFF)-oy)//16
            w=abs(x1-x0); h=abs(y1-y0)
            uv = uvbuf[-2:] if len(uvbuf)>=2 else [(0,0),(0,0)]
            xyzbuf=[]; uvbuf=[]
            ok = eligible(t0) and w<=MAX_SPRITE_PX and h<=MAX_SPRITE_PX and w>0 and h>0
            if not ok:
                if started: stop = "run ended: next sprite not v1-eligible"; break
                if t0 is None: rejected["no_tex0"]+=1
                elif not cur["tme"]: rejected["untextured"]+=1
                elif t0["psm"]!=PSMCT32: rejected[f"psm_0x{t0['psm']:02x}"]+=1
                elif not cur["fst"]: rejected["fst0_stq"]+=1
                elif not cur["abe"]: rejected["no_abe"]+=1
                elif t0["csm"] or t0["cpsm"]: rejected["clut"]+=1
                elif w>MAX_SPRITE_PX or h>MAX_SPRITE_PX: rejected["too_big"]+=1
                else: rejected["other"]+=1
                continue
            if not started:
                seg_tbp = t0["tbp"]; seg_t0 = t0; started=True; first=e.idx
            if t0["tbp"] != seg_tbp:
                stop = "run ended: TEX0 TBP changed (re-upload)"; break
            seg.append(dict(x0=min(x0,x1), y0=min(y0,y1), x1=max(x0,x1), y1=max(y0,y1), w=w, h=h,
                            u0=uv[0][0]/16.0, v0=uv[0][1]/16.0, u1=uv[1][0]/16.0, v1=uv[1][1]/16.0))
            last=e.idx
            if len(seg) >= max_spr: stop="hit --max"; break
    meta = dict(tbp=seg_tbp, tex0=(seg_t0 if started else None), first=first, last=last,
                stop=stop, rejected=dict(rejected))
    return seg, meta

def main(argv):
    if len(argv) < 2: print(__doc__); return 2
    path = argv[1]
    max_spr = int(argv[argv.index("--max")+1]) if "--max" in argv else DEF_MAX
    h, events = parse_dump(path)
    seg, meta = extract(events, max_spr)
    R = [f"SPRITE EXTRACT — serial={h.serial!r} crc=0x{h.crc:08x}  (v1: PSMCT32/UV/ABE, <= {MAX_SPRITE_PX}px)"]
    if not seg:
        R.append(f"NO v1-eligible sprite run selected. rejections: {meta['rejected']}")
    else:
        t0 = meta["tex0"]
        xs=[s['x0'] for s in seg]+[s['x1'] for s in seg]; ys=[s['y0'] for s in seg]+[s['y1'] for s in seg]
        sizes=Counter((s['w'],s['h']) for s in seg)
        R.append(f"selected {len(seg)} sprites, events #{meta['first']}..#{meta['last']}, stop: {meta['stop']}")
        R.append(f"  TEX0: TBP={t0['tbp']} TBW={t0['tbw']} {t0['tw']}x{t0['th']} PSMCT32 TFX={t0['tfx']} TCC={t0['tcc']}")
        R.append(f"  screen bbox: x[{min(xs)}..{max(xs)}] y[{min(ys)}..{max(ys)}]")
        R.append(f"  sizes WxH: {[(f'{w}x{hh}',n) for (w,hh),n in sizes.most_common(6)]}")
        u=[s['u0'] for s in seg]+[s['u1'] for s in seg]; vv=[s['v0'] for s in seg]+[s['v1'] for s in seg]
        R.append(f"  UV range: u[{min(u):.1f}..{max(u):.1f}] v[{min(vv):.1f}..{max(vv):.1f}] (texel coords)")
    rep = "\n".join(R)
    print(rep)
    if "--report" in argv:
        open(argv[argv.index("--report")+1], "w").write(rep+"\n")
    if "--out" in argv and seg:
        op = argv[argv.index("--out")+1]
        with open(op,"w") as f:
            t0=meta["tex0"]
            f.write(f"# Ch344 LOCAL authentic sprite run ({len(seg)}) — dump-derived. TBP={t0['tbp']} {t0['tw']}x{t0['th']} PSMCT32\n")
            f.write(f"tex0 {t0['tbp']} {t0['tbw']} {t0['tw']} {t0['th']} {t0['tfx']}\n")
            for s in seg:
                f.write(f"sprite {s['x0']} {s['y0']} {s['x1']} {s['y1']} {s['u0']:.1f} {s['v0']:.1f} {s['u1']:.1f} {s['v1']:.1f}\n")
        print(f"\nwrote sprite list -> {op}")
    return 0

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