#!/usr/bin/env python3
"""Plan a chronological, staging-safe SH3 PSMT8 coverage fixture.

This is the bridge from the draw census to gs_make_sh3_scheduler_fixture.py:
Select supported frame-1 textured draws, reconstruct texture/CLUT
identity in one local-memory replay, and group only consecutive draws whose
complete feeder state and sampled assets are identical.  The raw-triangle
word bound is conservative; the generator repeats the exact clipped check.
The default remains the historical opaque-strip plan.  Explicit switches can
also admit authentic ABE draws, triangle lists, and sprites for later fidelity
chapters.
"""
import argparse, hashlib, json, os, shlex, sys

HERE=os.path.dirname(os.path.abspath(__file__))
ROOT=os.path.normpath(os.path.join(HERE,".."))
sys.path.insert(0,HERE)
import gs_make_sh3_multidraw_fixture as MD
import gs_sh3_draw_census as C
import gs_sh3_recon as RC
import gs_texture_residency as R

STG_WORDS=2048
HEADER_WORDS=8

def normalized_state(e):
    s={k:e["state"].get(k) for k in ("prim","tex0","test","zbuf","clamp","alpha","texa","fogcol")}
    s=json.loads(json.dumps(s,sort_keys=True))
    t=s["tex0"]; clamp=s["clamp"]
    wms=clamp&3; wmt=(clamp>>2)&3
    minu=(clamp>>4)&0x3ff; maxu=(clamp>>14)&0x3ff
    minv=(clamp>>24)&0x3ff; maxv=(clamp>>34)&0x3ff
    if wms==2 and minu==0 and maxu==t["tw"]-1: clamp=(clamp&~3)|1
    if wmt==2 and minv==0 and maxv==t["th"]-1: clamp=(clamp&~12)|4
    s["clamp"]=clamp
    return json.dumps(s,sort_keys=True,separators=(",",":"))

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("dump",nargs="?")
    ap.add_argument("--census",default="/tmp/sh3_census.json")
    ap.add_argument("--frame",type=int,default=1,
                    help="capture frame to select (default: 1)")
    ap.add_argument("--tag",default="zsrt139f1")
    ap.add_argument("--out",default="/tmp/sh3_opaque_coverage_plan.json")
    ap.add_argument("--include-abe",action="store_true")
    ap.add_argument("--untextured-only",action="store_true",
                    help="select only TME=0 draws and emit the opt-in untextured feeder path")
    ap.add_argument("--prim-types",default="4",
                    help="comma-separated GS primitive types (3=triangles, 4=strip, 5=fan, 6=sprite)")
    ap.add_argument("--psms",default="0x13",
                    help="comma-separated texture PSMs (0x00=PSMCT32, 0x13=PSMT8, 0x14=PSMT4)")
    ap.add_argument("--min-idx",type=int)
    ap.add_argument("--max-idx",type=int)
    ap.add_argument("--native-fidelity",action="store_true",
                    help="emit the proven fog, bilinear, pixel-center, and native-12.4 flags")
    ap.add_argument("--fast-fit-scale",action="store_true",
                    help="choose the largest representable STQ scale instead of optimizing per-triangle UV error")
    ap.add_argument("--capacity-epoch-pixels",type=int,default=14500,
                    help="software coverage cap used by the 16-row capacity splitter")
    ap.add_argument("--chapter",default="Ch409")
    ns=ap.parse_args()
    dump=ns.dump
    if not dump:
        import glob
        hits=glob.glob(os.path.join(ROOT,"captures","gs","silenthill3","*224139*.gs.zst"))
        if not hits: sys.exit("no 224139 dump found")
        dump=hits[0]

    collected=R.collect(dump,0)
    if os.path.exists(ns.census) and not ns.untextured_only:
        census=json.load(open(ns.census))
    else:
        census,_,_=C.census(dump,frame_filter=ns.frame,collected=collected,
                            include_untextured=ns.untextured_only)
    prim_types={int(x) for x in ns.prim_types.split(",")}
    if not prim_types or not prim_types.issubset({3,4,5,6}):
        sys.exit(f"--prim-types must be a nonempty subset of 3,4,5,6; got {sorted(prim_types)}")
    psms={int(x,0) for x in ns.psms.split(",")}
    if not psms or not psms.issubset({0x00,0x13,0x14}):
        sys.exit(f"--psms must be a nonempty subset of 0x00,0x13,0x14; got {sorted(psms)}")
    def supported_texture(d):
        if ns.untextured_only:
            return True
        t=d["tex0"]
        return (t["psm"] in psms and
                ((t["psm"]==0x00 and t["tw"]==64 and t["th"]==64 and t["tbw"]==1)
                 or (t["psm"]==0x13 and t["tw"] in (128,256,512) and t["th"] in (128,256,512))
                 or (t["psm"]==0x14 and t["tw"]==512 and t["th"]==1024)))
    selected=sorted(d["first_idx"] for d in census
                    if d["frame"]==ns.frame and d.get("tex_resident")
                    and d["prim"]["type"] in prim_types
                    and d["prim"]["tme"]==(0 if ns.untextured_only else 1)
                    and (ns.untextured_only or d["prim"]["fst"]==0)
                    and d["prim"]["ctxt"]==0
                    and (ns.include_abe or d["prim"]["abe"]==0)
                    and (ns.min_idx is None or d["first_idx"]>=ns.min_idx)
                    and (ns.max_idx is None or d["first_idx"]<=ns.max_idx)
                    and supported_texture(d))
    got,_=MD.load_draws(dump,selected,collected=collected,
                        allow_untextured=ns.untextured_only)
    missing=sorted(set(selected)-set(got))
    if missing: sys.exit(f"capture lost selected draws: {missing[:10]}")

    rows=[]
    for idx,mem in RC.iter_localmem_at(dump,selected,collected=collected):
        e=got[idx]; t=e["state"]["tex0"]
        if ns.untextured_only:
            tex=bytes(512*512); clut=bytes(1024)
        elif t["psm"]==0x00:
            tex=b"".join(mem.read_ct32_word(t["tbp"],t["tbw"],x,y).to_bytes(4,"little")
                         for y in range(t["th"]) for x in range(t["tw"]))
        elif t["psm"]==0x13:
            tex=bytes(mem.read_psmt8(t["tbp"],t["tbw"],512,512))
        else:
            unpacked=mem.read_psmt4(t["tbp"],t["tbw"],t["tw"],t["th"])
            tex=bytes((unpacked[n]&15)|((unpacked[n+1]&15)<<4)
                      for n in range(0,len(unpacked),2))
        if not ns.untextured_only:
            clut=(b"" if t["psm"]==0x00 else
                  bytes(mem.m[t["cbp"]*256:t["cbp"]*256+1024]))
        asset=hashlib.sha256(tex+clut).hexdigest()
        ptype=e["state"]["prim"]["type"]
        if ptype==3:
            ntri=len(e["verts"])//3
        elif ptype==4:
            ntri=max(0,len(e["verts"])-2)
        elif ptype==5:
            ntri=max(0,len(e["verts"])-2)
        else:
            # One GS sprite is a pair of opposing corners and expands into
            # two independent triangles in the scheduler fixture.
            ntri=len(e["verts"])
        words=HEADER_WORDS+9*ntri
        if words>STG_WORDS:
            sys.exit(f"idx{idx}: single draw raw bound {words}>{STG_WORDS}")
        rows.append(dict(idx=idx,ntri=ntri,state=normalized_state(e),asset=asset))

    groups=[]; cur=[]; cur_tris=0; cur_key=None
    for r in rows:
        key=(r["state"],r["asset"])
        if cur and (key!=cur_key or HEADER_WORDS+9*(cur_tris+r["ntri"])>STG_WORDS):
            groups.append(cur); cur=[]; cur_tris=0
        cur.append(r); cur_tris+=r["ntri"]; cur_key=key
    if cur: groups.append(cur)

    idxarg=",".join(str(r["idx"]) for r in rows)
    grouparg=",".join(str(len(g)) for g in groups)
    cmd=["python3","tools/gs_make_sh3_scheduler_fixture.py",dump,
         "--draw-list",idxarg,"--group-sizes",grouparg,"--authz","--fb640",
         "--tag",ns.tag,"--runtime-clut","--allow-abe","--auth-color-tfx",
         "--emit-clamp-header","--normalize-full-region-clamp","--clip-fb-guardband","--auth-scissor-clip",
         "--capacity-epochs",str(ns.capacity_epoch_pixels),
         "--pscale",("1" if ns.untextured_only else "auto"),
         "--skip-mixed-q-tris"]
    if ns.fast_fit_scale:
        cmd += ["--fast-fit-scale"]
    if ns.untextured_only:
        cmd += ["--allow-untextured","--coalesce-solid-sprites","--half-open-sprites"]
    if ns.native_fidelity:
        cmd += ["--auth-fog-black-fold","--ref-bilinear","--ref-sample","center","--subpixel-xy"]
    cmd += ["--chapter",ns.chapter,"--emit"]
    out=dict(dump=dump,tag=ns.tag,n_draws=len(rows),n_groups=len(groups),
             max_group_draws=max(map(len,groups)),max_group_raw_words=max(HEADER_WORDS+9*sum(r["ntri"] for r in g) for g in groups),
             include_abe=ns.include_abe,prim_types=sorted(prim_types),psms=sorted(psms),
             untextured_only=ns.untextured_only,
             frame=ns.frame,
             min_idx=ns.min_idx,max_idx=ns.max_idx,
             native_fidelity=ns.native_fidelity,
             capacity_epoch_pixels=ns.capacity_epoch_pixels,
             draw_ids=[r["idx"] for r in rows],group_sizes=[len(g) for g in groups],command=cmd)
    with open(ns.out,"w") as f: json.dump(out,f,indent=2); f.write("\n")
    mode="opaque" if not ns.include_abe else "opaque+ABE"
    print(f"[{ns.chapter}] {len(rows)} {mode} indexed PSM {sorted(hex(x) for x in psms)} type {sorted(prim_types)} draws -> {len(groups)} exact-state/asset groups")
    print(f"[{ns.chapter}] max group {out['max_group_draws']} draws; conservative max {out['max_group_raw_words']}/{STG_WORDS} words")
    print(f"[{ns.chapter}] plan {ns.out}")
    print(shlex.join(cmd))

if __name__=="__main__":
    main()
