#!/usr/bin/env python3
"""retroDE_ps2 — SPRITE census (Ch344 sprite-ingestion scoping).

Consumes gs_parse.parse_dump()'s event stream and characterises every SPRITE primitive in a .gs dump:
prim-type split, per-sprite TEX0 PSM, TME/ABE/FST, the ALPHA blend equation, TEX1 mag/min filter, and an
XYOFFSET-corrected size histogram. Output is AGGREGATE (counts/histograms) only — no copyrighted pixel
content — so the report is committable per captures/gs/.gitignore policy. This is census/scoping ONLY;
it renders nothing and asserts nothing.

Usage:  gs_census_sprites.py <dump.gs[.xz|.zst]> [--report out.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

PSM = {0x00:"PSMCT32",0x01:"PSMCT24",0x02:"PSMCT16",0x0A:"PSMCT16S",0x13:"PSMT8",0x14:"PSMT4",
       0x1B:"PSMT8H",0x24:"PSMT4HL",0x2C:"PSMT4HH"}
PRIMT = {0:"POINT",1:"LINE",2:"LINE_STRIP",3:"TRI",4:"TRISTRIP",5:"TRIFAN",6:"SPRITE",7:"INVALID"}
ABCD = {0:"Cs",1:"Cd",2:"0"}; C_SEL = {0:"As",1:"Ad",2:"FIX"}
FILT = {0:"NEAREST",1:"LINEAR",2:"N_MIPN",3:"L_MIPN",4:"N_MIPL",5:"L_MIPL"}

def dec_tex0(v):
    return (v>>20)&0x3F, 1<<((v>>26)&0xF), 1<<((v>>30)&0xF), v&0x3FFF, (v>>35)&3, (v>>34)&1
def dec_alpha(v):
    a=v&3; b=(v>>2)&3; c=(v>>4)&3; d=(v>>6)&3; fix=(v>>32)&0xFF
    return f"({ABCD[a]}-{ABCD[b]})*{C_SEL[c]}>>7+{ABCD[d]}" + (f" FIX={fix}" if c==2 else "")

def census(path):
    h, events = parse_dump(path)
    cur = {"type":None,"tme":0,"abe":0,"fst":0}
    tex0 = {1:None, 2:None}; xyoff = {1:(0,0), 2:(0,0)}; tex1 = {1:None, 2:None}
    alpha = {1:None, 2:None}; ctxt = 0
    kicks = {}; xyz = []
    spr = {"n":0, "psm":Counter(), "abe":Counter(), "fst":Counter(), "tme":Counter(),
           "alpha":Counter(), "magfilt":Counter(), "sizes":Counter()}
    for e in events:
        if e.kind != "GSREG": continue
        a, v = e.addr, e.value
        if a == 0x00:                                   # PRIM
            cur = {"type":v&7,"tme":(v>>4)&1,"abe":(v>>6)&1,"fst":(v>>8)&1}; ctxt=(v>>9)&1; xyz=[]
        elif a in (0x06,0x07): tex0[1 if a==0x06 else 2] = dec_tex0(v)
        elif a in (0x14,0x15): tex1[1 if a==0x14 else 2] = (v>>5)&7    # MMAG bit5 -> mag filter (1bit here: 0/1)
        elif a in (0x18,0x19): xyoff[1 if a==0x18 else 2] = (v&0xFFFF, (v>>32)&0xFFFF)
        elif a in (0x42,0x43): alpha[1 if a==0x42 else 2] = v
        elif a == 0x05:                                  # XYZ2 vertex kick
            t = cur["type"]
            if t is None: continue
            kicks[t] = kicks.get(t,0)+1
            if t == 6:
                xyz.append(v)
                if len(xyz) == 2:
                    cx = 1 if ctxt==0 else 2
                    ox, oy = xyoff[cx]
                    x0=(xyz[0]&0xFFFF)-ox; y0=((xyz[0]>>16)&0xFFFF)-oy
                    x1=(xyz[1]&0xFFFF)-ox; y1=((xyz[1]>>16)&0xFFFF)-oy
                    w=abs(x1-x0)//16; ht=abs(y1-y0)//16            # 12.4 fixed -> pixels
                    spr["n"]+=1; spr["sizes"][(w,ht)]+=1
                    spr["abe"][cur["abe"]]+=1; spr["fst"][cur["fst"]]+=1; spr["tme"][cur["tme"]]+=1
                    tx = tex0[cx]
                    spr["psm"][("(untextured)" if not cur["tme"] else
                                "(no TEX0)" if tx is None else PSM.get(tx[0],f"0x{tx[0]:02x}"))] += 1
                    if cur["abe"]:
                        al = alpha[cx]
                        spr["alpha"][dec_alpha(al) if al is not None else "(no ALPHA set)"] += 1
                    if cur["tme"]:
                        f1 = tex1[cx]
                        spr["magfilt"][("NEAREST" if f1==0 else "LINEAR" if f1==1 else "(unset)") ] += 1
                    xyz=[]
    return h, kicks, spr

def fmt(h, kicks, spr):
    L = []
    L.append(f"SPRITE CENSUS — serial={h.serial!r} crc=0x{h.crc:08x}")
    L.append(f"prim-kicks by type: {dict((PRIMT[k],v) for k,v in sorted(kicks.items(),key=lambda x:-x[1]))}")
    L.append(f"SPRITES: {spr['n']} rectangles")
    L.append(f"  TEX0 PSM      : {dict(spr['psm'].most_common())}")
    L.append(f"  TME(textured) : {dict(spr['tme'])}   FST(0=STQ,1=UV): {dict(spr['fst'])}")
    L.append(f"  ABE(alpha)    : {dict(spr['abe'])}")
    L.append(f"  ALPHA eqn     : {dict(spr['alpha'].most_common())}")
    L.append(f"  mag filter    : {dict(spr['magfilt'].most_common())}")
    L.append(f"  top sizes WxH : {[(f'{w}x{ht}',n) for (w,ht),n in spr['sizes'].most_common(8)]}")
    return "\n".join(L)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__); sys.exit(2)
    h, kicks, spr = census(sys.argv[1])
    out = fmt(h, kicks, spr)
    print(out)
    if "--report" in sys.argv:
        p = sys.argv[sys.argv.index("--report")+1]
        with open(p, "w") as f: f.write(out + "\n")
        print(f"\nwrote report -> {p}")