#!/usr/bin/env python3
"""retroDE_ps2 — Ch340 support census + translator (Bricks 3-4).

Consumes gs_parse's normalized event stream, reconstructs GS primitives via the vertex-kick model,
classifies EVERY primitive (translated / unsupported, with reason + frame/event/byte offset), emits an
aggregate census + histograms, then translates the EARLIEST mechanically-selected contiguous SUPPORTED
draw segment into a ps2_feeder scene file (Ch339 grammar) — reusing the proven encoder, never
duplicating staging logic. NO hidden approximation: textured prims, sprites, strips/fans, non-source-
over blend, unsupported Z-test, etc. are reported unsupported and the segment STOPS there (fail closed).
If no segment qualifies, Ch340 still succeeds via parser + census and the top census blocker frames Ch341.

A declared, reported VIEWPORT FIT (game bbox -> our 64x64 FB) is a faithful linear transform, not a GS-
feature approximation; it is recorded in the qualification report.

Usage:
  gs_translate.py <dump.gs[.xz|.zst]> [--report out.census.txt] [--scene out.scene.txt] [--fb N]
"""
import sys, os
sys.path.insert(0, os.path.dirname(__file__))
import gs_parse

PRIMT = {0:"POINT",1:"LINE",2:"LINE_STRIP",3:"TRIANGLE",4:"TRI_STRIP",5:"TRI_FAN",6:"SPRITE",7:"INVALID"}
def bit(v,b): return (v>>b)&1

class Prim:
    __slots__=("type","verts","tme","abe","iip","frame","event_idx","byte_off","test","alpha")
    def __init__(s,**k):
        for a in s.__slots__: setattr(s,a,k.get(a))

def reconstruct_prims(events):
    """Vertex-kick model -> list of Prim. Each vertex = (x,y,z,(r,g,b)) in raw GS coords."""
    prims=[]
    st=dict(prim=None, rgbaq=0, xyoff=0, tme=0, abe=0, iip=0, test=0, alpha=0, frame=0, ptype=7)
    vtx=[]
    for e in events:
        if e.kind!="GSREG":
            if e.kind=="FRAME_BOUNDARY": vtx=[]   # vertex FIFO doesn't survive a frame in our model
            continue
        r=e.reg; v=e.value
        if r=="PRIM":
            st["prim"]=v; st["ptype"]=v&7; st["iip"]=bit(v,3); st["tme"]=bit(v,4); st["abe"]=bit(v,6)
            vtx=[]
        elif r=="PRMODE":
            st["iip"]=bit(v,3); st["tme"]=bit(v,4); st["abe"]=bit(v,6)
        elif r=="RGBAQ": st["rgbaq"]=v
        elif r=="XYOFFSET_1": st["xyoff"]=v
        elif r in ("TEST_1",): st["test"]=v
        elif r in ("ALPHA_1",): st["alpha"]=v
        elif r in ("FRAME_1",): st["frame"]=v
        elif r in ("XYZ2","XYZ3","XYZF2","XYZF3"):
            x=(v&0xFFFF); y=((v>>16)&0xFFFF); z=(v>>32)&0xFFFFFFFF
            col=(st["rgbaq"]&0xFF,(st["rgbaq"]>>8)&0xFF,(st["rgbaq"]>>16)&0xFF)
            kick = r in ("XYZ2","XYZF2")    # XYZ3/XYZF3 add a vertex without a drawing kick
            vtx.append((x,y,z,col))
            t=st["ptype"]
            need = {0:1,1:2,2:2,3:3,4:3,5:3,6:2}.get(t,99)
            if kick and len(vtx)>=need:
                pv = vtx[-need:]
                prims.append(Prim(type=t,verts=pv,tme=st["tme"],abe=st["abe"],iip=st["iip"],
                                  frame=e.frame,event_idx=e.idx,byte_off=e.byte_off,test=st["test"],alpha=st["alpha"]))
                if t==3 or t==6: vtx=[]            # independent tri / sprite: consume the FIFO
                # strips/fans keep a sliding window (left as-is; classified unsupported below)
    return prims, st

# ---- proven envelope: a primitive we can render faithfully via ps2_feeder ----
def classify(p):
    if p.type!=3:           return "unsupported", f"prim={PRIMT.get(p.type,p.type)} (only TRIANGLE renders via ps2_feeder)"
    if p.tme:               return "unsupported", "textured triangle (TME=1; no real-texture path in the feeder)"
    if p.abe:               return "unsupported", "alpha-blended triangle (ABE=1; only opaque is in the proven envelope)"
    return "translated", "non-textured opaque triangle"

def census(prims, parse_summary):
    cats={}; reasons={}
    for p in prims:
        c,why=classify(p)
        cats[c]=cats.get(c,0)+1
        if c=="unsupported": reasons[why]=reasons.get(why,0)+1
    return cats, reasons

def earliest_supported_segment(prims):
    """Earliest maximal contiguous run of 'translated' prims (stops at the first unsupported)."""
    seg=[]; best=None
    for p in prims:
        c,_=classify(p)
        if c=="translated":
            seg.append(p)
        else:
            if seg: best=seg; break          # earliest contiguous run -> stop at first unsupported AFTER it
            seg=[]
    if not best and seg: best=seg
    return best or []

def viewport_fit(prims, fb):
    """Declared linear map of the segment's PIXEL bbox into [1, fb-2] (margin), reported. f works on
    RAW (1/16) GS coords; bbox + scale are in screen PIXELS so the report reads intuitively."""
    xs=[x/16.0 for p in prims for (x,y,z,c) in p.verts]; ys=[y/16.0 for p in prims for (x,y,z,c) in p.verts]
    x0,x1,y0,y1=min(xs),max(xs),min(ys),max(ys)
    s=min((fb-2)/max(1.0,(x1-x0)), (fb-2)/max(1.0,(y1-y0)))      # output px per source px
    def f(x,y): return (1+int((x/16.0-x0)*s), 1+int((y/16.0-y0)*s))
    return f, dict(bbox_px=(round(x0,1),round(y0,1),round(x1,1),round(y1,1)), scale_px=round(s,4), fb=fb)

def emit_scene(seg, fb):
    f,info=viewport_fit(seg, fb)
    lines=[f"# Ch340 translated segment: {len(seg)} non-textured triangles",
           f"# viewport fit: src bbox(px) {info['bbox_px']} scale {info['scale_px']} px/px -> {fb}x{fb}"]
    for p in seg:
        pts=[]
        for (x,y,z,c) in p.verts:
            X,Y=f(x,y); pts += [X,Y]
        (r0,g0,b0)=p.verts[0][3]
        if p.iip:   # gouraud: trig x0 y0 r0 g0 b0  x1 y1 ...  z
            a=[]
            for (x,y,z,c) in p.verts:
                X,Y=f(x,y); a += [X,Y,c[0],c[1],c[2]]
            z=p.verts[-1][2]
            lines.append("trig "+" ".join(map(str,a+[z])))
        else:       # flat: tri x0 y0 x1 y1 x2 y2 z r g b   (GS flat uses provoking/last vertex color)
            (r,g,b)=p.verts[-1][3]; z=p.verts[-1][2]
            lines.append(f"tri {pts[0]} {pts[1]} {pts[2]} {pts[3]} {pts[4]} {pts[5]} {z} {r} {g} {b}")
    lines.append("go")
    return "\n".join(lines)+"\n", info

def main(argv):
    if len(argv)<2: print(__doc__); return 2
    path=argv[1]
    def opt(name,d=None):
        return argv[argv.index(name)+1] if name in argv else d
    fb=int(opt("--fb","64"))
    h, events = gs_parse.parse_dump(path)
    prims, _ = reconstruct_prims(events)
    cats, reasons = census(prims, h)

    # histograms (from the raw event stream)
    regs={}; flgs={}; frames=0; images=0; imgb=0; malformed=0
    ptypes={}
    for e in events:
        if e.kind=="GSREG": regs[e.reg]=regs.get(e.reg,0)+1
        if e.kind=="FRAME_BOUNDARY": frames+=1
        if e.kind=="MALFORMED": malformed+=1
        if e.kind=="IMAGE": images+=1; imgb+=e.info.get("bytes",0)
    for p in prims: ptypes[PRIMT.get(p.type,p.type)]=ptypes.get(PRIMT.get(p.type,p.type),0)+1

    R=[]
    R.append(f"# Ch340 GS-dump support census  (schema v{gs_parse.SCHEMA_VERSION})")
    R.append(f"# source: {os.path.basename(path)}  serial={h.serial!r} crc=0x{h.crc:08x}  (aggregate counts only; no game content)")
    R.append(f"frames={frames}  events={len(events)}  primitives={len(prims)}  malformed={malformed}  image_uploads={images} ({imgb} bytes)")
    R.append("")
    R.append("primitive types: "+", ".join(f"{k}={v}" for k,v in sorted(ptypes.items(),key=lambda x:-x[1])))
    R.append("census classes : "+", ".join(f"{k}={v}" for k,v in sorted(cats.items(),key=lambda x:-x[1])))
    R.append("")
    R.append("UNSUPPORTED reasons (count):")
    for why,c in sorted(reasons.items(),key=lambda x:-x[1]): R.append(f"  {c:6d}  {why}")
    R.append("")
    R.append("top GS register writes:")
    for k,v in sorted(regs.items(),key=lambda x:-x[1])[:18]: R.append(f"  {v:6d}  {k}")

    seg=earliest_supported_segment(prims)
    R.append("")
    if seg:
        first=seg[0]
        R.append(f"EARLIEST SUPPORTED SEGMENT: {len(seg)} triangles, starting frame {first.frame} "
                 f"event #{first.event_idx} @0x{first.byte_off:x}.")
        R.append("  qualification: every primitive is a non-textured opaque TRIANGLE; segment stops at the first unsupported prim/state.")
    else:
        topblk = max(reasons.items(), key=lambda x:x[1])[0] if reasons else "none"
        R.append("NO SUPPORTED SEGMENT: no contiguous run of non-textured opaque triangles.")
        R.append(f"  Ch340 succeeds via parser + census. Top census blocker (Ch341 candidate): {topblk}")

    report="\n".join(R)+"\n"
    print(report)
    rp=opt("--report")
    if rp:
        with open(rp,"w") as f: f.write(report)
        print(f"[wrote census -> {rp}]")
    sp=opt("--scene")
    if sp and seg:
        scene,info=emit_scene(seg,fb)
        with open(sp,"w") as f: f.write(scene)
        print(f"[wrote {len(seg)}-tri ps2_feeder scene -> {sp}  (viewport {info})]")
    elif sp:
        print("[no scene emitted: no supported segment — fail closed]")
    return 0

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