#!/usr/bin/env python3
"""Emit the authentic SH3 4x4 PSMT4 post-process tile pass as scheduler epochs.

The dump expresses this pass as sixteen 128x128 FST SPRITEs.  The runtime
feeder is triangle-record based, so each sprite is expanded to two textured
triangles and split into ordered 16-row bands.  Consecutive bands are packed
under the proven 14,500-fragment request-FIFO safety bound.  The PSMT4 asset
is de-swizzled to the texture cache's linear packed-nibble layout and padded
to its fixed 256 KiB residency window.
"""
import glob
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.normpath(os.path.join(HERE, ".."))
DATA = os.path.join(ROOT, "sim", "data", "top_psmct32_raster_demo")
sys.path.insert(0, HERE)
sys.path.insert(0, DATA)
import gs_make_sh3_multidraw_fixture as MD
import gs_sh3_recon as RC
import bake

DRAW_IDS = [198945,198954,198963,198972,198981,198990,198999,199008,
            199017,199026,199035,199044,199053,199062,199071,199080]
FBW, W, H = 10, 640, 480
NEW_TBP, CBP = 1024, 480
TEX_CACHE_BYTES = 512*512
STG_WORDS = 2048
CAPACITY = 14500
PERSP_FRAC = bake.PERSP_FRAC
PSCALE = 1/64


def wmem(path, words, banner, width=8):
    with open(os.path.join(DATA, path), "w") as f:
        f.write(f"// {banner}\n")
        for word in words:
            f.write(f"{word & ((1 << (4*width))-1):0{width}x}\n")


def vwords(v):
    s_fp = round(v["s"] * 128 * (1 << PERSP_FRAC) * PSCALE)
    t_fp = round(v["t"] * 128 * (1 << PERSP_FRAC) * PSCALE)
    q_fp = round((1 << PERSP_FRAC) * PSCALE)
    x = max(0, min(W-1, round(v["x"])))
    y = max(0, min(H-1, round(v["y"])))
    return [bake.rgbaq_with_q(0x80, 0x80, 0x80, q_fp),
            bake.st_data(s_fp, t_fp), bake.xyz2_dataz(x, y, 16)]


def main(argv):
    tag = argv[argv.index("--tag")+1] if "--tag" in argv else "zsrt139e1"
    dump = next((x for x in argv[1:] if x.endswith(".gs.zst")), None)
    if dump is None:
        found = glob.glob(os.path.join(ROOT, "captures", "gs", "silenthill3", "*224139*.gs.zst"))
        if not found:
            sys.exit("no dump 224139 found")
        dump = found[0]

    got, _ = MD.load_draws(dump, DRAW_IDS)
    if sorted(got) != DRAW_IDS:
        sys.exit(f"missing overlay draws: {sorted(set(DRAW_IDS)-set(got))}")
    states = [got[i]["state"] for i in DRAW_IDS]
    ref = states[0]
    for i, state in zip(DRAW_IDS, states):
        t0, pr = state["tex0"], state["prim"]
        if (t0["tbp"],t0["tbw"],t0["psm"],t0["tw"],t0["th"],t0["cbp"]) != (13376,2,0x14,128,128,14276):
            sys.exit(f"idx{i}: unexpected TEX0 {t0}")
        if (pr["type"],pr["tme"],pr["fst"],pr["abe"]) != (6,1,1,1):
            sys.exit(f"idx{i}: unexpected PRIM {pr}")
        for key in ("test","zbuf","clamp","scissor","alpha","texa"):
            if state[key] != ref[key]:
                sys.exit(f"idx{i}: {key} differs")
    if ref["scissor"] != 0x01ff000001ff0000:
        sys.exit(f"unexpected SCISSOR 0x{ref['scissor']:016x}")

    mem, *_ = RC.build_localmem_to(dump, DRAW_IDS[0])
    idx = mem.read_psmt4(13376, 2, 128, 128)
    packed = bytearray(128*128//2)
    for n in range(0, len(idx), 2):
        packed[n//2] = (idx[n]&15) | ((idx[n+1]&15)<<4)
    packed.extend(bytes(TEX_CACHE_BYTES-len(packed)))
    tex_words = [int.from_bytes(packed[n:n+4], "little") for n in range(0,TEX_CACHE_BYTES,4)]
    crc = sum(tex_words) & 0xFFFFFFFF
    pal = RC.read_clut32(mem, 14276, order="grid")

    # Convert each authentic sprite to 16-row triangle bands.  The sampled V
    # coordinate is intentionally constant zero in the dump; U spans 0..128.
    bands = []
    for draw_id in DRAW_IDS:
        vv = got[draw_id]["verts"]
        if len(vv) != 2:
            sys.exit(f"idx{draw_id}: expected two sprite vertices, got {len(vv)}")
        a,b = vv
        x0,x1 = a["x"],b["x"]
        y0,y1 = a["y"],min(b["y"],float(H))
        if y1 <= 0 or y0 >= H:
            continue
        for lo in range(max(0,int(round(y0))), int(round(y1)), 16):
            hi = min(int(round(y1)), lo+16)
            tl=dict(x=x0,y=lo,s=a["s"],t=a["t"])
            tr=dict(x=x1,y=lo,s=b["s"],t=a["t"])
            bl=dict(x=x0,y=hi,s=a["s"],t=b["t"])
            br=dict(x=x1,y=hi,s=b["s"],t=b["t"])
            cov=max(0,round(x1)-round(x0))*max(0,hi-lo)
            bands.append((draw_id,[(tl,tr,bl),(tr,br,bl)],cov))

    epochs=[]; cur=[]; cov=0
    for draw_id,tris,npx in bands:
        if cur and cov+npx > CAPACITY:
            epochs.append((cur,cov)); cur=[]; cov=0
        cur.extend(tris); cov += npx
    if cur:
        epochs.append((cur,cov))
    if max(x[1] for x in epochs) > CAPACITY:
        sys.exit("capacity pack failure")

    tex0 = bake.tex0_pack(NEW_TBP,2,psm=0x14,tw=7,th=7,tfx=0)
    tex0 |= (CBP&0x3FFF)<<37  # runtime CLUT: CLD=0
    prim = 3 | (1<<4) | (1<<6)  # TRIANGLE, TME, ABE; FST converted to ST/Q
    for k,(tris,npx) in enumerate(epochs):
        stg=[len(tris)|(1<<32)|(1<<34), bake.frame_1_psmct32(FBW), ref["alpha"],
             ref["test"], ref["zbuf"], tex0, ref["clamp"], prim]
        for tri in tris:
            for v in tri:
                stg.extend(vwords(v))
        if len(stg)>STG_WORDS:
            sys.exit(f"epoch{k}: staging overflow {len(stg)}")
        bake.write_feeder_stg_mem(f"feeder_sh3_{tag}{k}.mem",stg,
            f"Ch408 authentic PSMT4 overlay epoch{k}, {len(tris)} tris, capacity {npx}. gitignored.",total=STG_WORDS)
        wmem(f"sh3_{tag}{k}_tex_lpddr.mem",tex_words,
             f"Ch408 PSMT4 packed-linear texture padded to fixed 256 KiB cache; crc=0x{crc:08x}. gitignored.")
        wmem(f"sh3_{tag}{k}_pal.mem",pal,
             f"Ch408 authentic PSMT4 CSM1 palette, runtime staged. gitignored.")

    n=len(epochs)
    with open(os.path.join(DATA,f"sh3_{tag}_epochs.txt"),"w") as f:
        f.write("# Ch408 authentic PSMT4 post-process overlay scheduler fixture\n")
        f.write("# k idx tbp cbp_reloc tex_file lpddr size crc list_file words records reuse pal_file pal_sum32\n")
        f.write(f"META n_epochs {n} fbpxw {W} fbh {H} fbwords {W*H} lpddr_tex 0x200000 tex_words {TEX_CACHE_BYTES//4} n_beats {TEX_CACHE_BYTES//32}\n")
        for k,(tris,npx) in enumerate(epochs):
            words=8+9*len(tris)
            f.write(f"{k} {DRAW_IDS[0]} 13376 {CBP} sh3_{tag}{k}_tex_lpddr.mem 0x200000 {TEX_CACHE_BYTES} 0x{crc:08x} "
                    f"feeder_sh3_{tag}{k}.mem {words} {len(tris)} {0 if k==0 else 1} sh3_{tag}{k}_pal.mem 0x{sum(pal)&0xFFFFFFFF:08x}\n")
    with open(os.path.join(DATA,f"sh3_{tag}_params.vh"),"w") as f:
        f.write("// Ch408 generated params. gitignored.\n")
        for key,val in (("FBW",FBW),("FBPXW",W),("FBH",H),("VRAM_BYTES_P",128*1024),("NEW_TBP",NEW_TBP),
                        ("TEX_VRAM_BASE",NEW_TBP*256),("TEX_BYTES",TEX_CACHE_BYTES),("N_BEATS",TEX_CACHE_BYTES//32),
                        ("STG_WORDS",STG_WORDS),("TW",128),("TH",128),("N_EPOCHS",n),("UNION_OX",0),("UNION_OY",0)):
            f.write(f"localparam int {key:<14}= {val};\n")
        f.write("localparam [29:0] LPDDR_TEX_BASE = 30'h0200000;\n")
        for k,(tris,npx) in enumerate(epochs):
            f.write(f"localparam int EP{k}_NTRIS={len(tris)}; localparam int EP{k}_CBP={CBP}; "
                    f"localparam [31:0] EP{k}_CRC=32'h{crc:08x}; localparam bit EP{k}_REUSE=1'b{0 if k==0 else 1};\n")
    with open(os.path.join(DATA,f"sh3_{tag}_epoch_table.vh"),"w") as f:
        for k in range(n):
            f.write(f"EP_CRC[{k}]=EP{k}_CRC; EP_REC[{k}]=EP{k}_NTRIS; EP_REUSE[{k}]=EP{k}_REUSE;\n")
    # Runtime CLUT mode uses an empty boot payload.
    qwc=0; disp=((H-1)<<12)|(W-1)
    with open(os.path.join(DATA,f"payload_sh3_{tag}.mem"),"w") as f:
        f.write(f"// Ch408 runtime-CLUT empty payload (CBPs [{CBP}]). gitignored. QWC={qwc}.\n")
        for _ in range(512): f.write("00000000000000000000000000000000\n")
    bake.write_bios_mem(f"bios_sh3_{tag}.mem",bake.build_textured_demo_bootlet_disp(qwc,disp,FBW),
                        f"Ch408 runtime PSMT4 overlay bootlet DISPLAY1={W}x{H}. gitignored.")
    print(f"[Ch408] PASS: {len(DRAW_IDS)} sprites -> {len(bands)} bands -> {n} epochs; max coverage={max(x[1] for x in epochs)}")
    print(f"[Ch408] PSMT4 indices={sorted(set(idx))}; texture cache crc=0x{crc:08x}; palette sum=0x{sum(pal)&0xFFFFFFFF:08x}")
    return 0


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