#!/usr/bin/env python3
"""Emit SH3's two authentic PSMCT32 full-frame darkening sprites.

The source frame builds a 256x256 direct-color intermediate at TBP 11264,
then draws it over the 512-wide display as two ABE sprites (indices 196167
and 196177).  The production texture cache is exactly 256 KiB, so the
intermediate can be reconstructed from GS local memory, linearized, and
replayed without first implementing the two off-screen render targets.
"""
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 = (196167, 196177)
W, H, FBW = 640, 480, 10
TW = TH = 256
NEW_TBP = 1024
TEX_BYTES = TW * TH * 4
N_BEATS = TEX_BYTES // 32
STG_WORDS = 2048
CAPACITY = 14500
PERSP_FRAC = bake.PERSP_FRAC
PSCALE = 4


def wmem(name, words, banner):
    with open(os.path.join(DATA, name), "w") as f:
        f.write(f"// {banner}\n")
        for word in words:
            f.write(f"{word & 0xffffffff:08x}\n")


def vertex_words(v):
    s_fp = round(v["s"] * TW * (1 << PERSP_FRAC) * PSCALE)
    t_fp = round(v["t"] * TH * (1 << PERSP_FRAC) * PSCALE)
    q_fp = round(v["q"] * (1 << PERSP_FRAC) * PSCALE)
    if not (-(1 << 23) <= s_fp < (1 << 23) and -(1 << 23) <= t_fp < (1 << 23)):
        raise ValueError(f"S/T overflow: {s_fp},{t_fp}")
    x = max(0, min(W - 1, round(v["x"])))
    y = max(0, min(H - 1, round(v["y"])))
    rgba = v["rgba"]
    # Preserve the authentic RGBAQ alpha.  bake.rgbaq_with_q() is a legacy
    # opaque helper (A=0xff), but these MODULATE+TCC passes depend on the
    # captured 0x50 vertex alpha: As = At*Av/128.
    return [((q_fp & 0xffffffff) << 32) | (rgba & 0xffffffff),
            bake.st_data(s_fp & 0xffffff, t_fp & 0xffffff),
            bake.xyz2_dataz(x, y, v["z"])]


def sprite_bands(draw):
    a, b = draw["verts"]
    if not (a.get("kick", True) and b.get("kick", True)):
        return []
    y0 = max(0.0, a["y"])
    y1 = min(float(H), b["y"])
    if y1 <= y0:
        return []
    out = []
    for lo in range(int(y0), int(y1), 16):
        hi = min(int(y1), lo + 16)
        fa = (lo - a["y"]) / (b["y"] - a["y"])
        fb = (hi - a["y"]) / (b["y"] - a["y"])
        top = dict(a, y=float(lo), t=a["t"] + fa * (b["t"] - a["t"]),
                   q=a["q"] + fa * (b["q"] - a["q"]))
        bot = dict(b, y=float(hi), t=a["t"] + fb * (b["t"] - a["t"]),
                   q=a["q"] + fb * (b["q"] - a["q"]))
        tl = dict(top, x=a["x"], s=a["s"])
        tr = dict(top, x=b["x"], s=b["s"])
        bl = dict(bot, x=a["x"], s=a["s"])
        br = dict(bot, x=b["x"], s=b["s"])
        tris = ((tl, tr, bl), (tr, br, bl))
        coverage = max(0, round(b["x"]) - round(a["x"])) * (hi - lo)
        out.append((tris, coverage))
    return out


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

    got, _ = MD.load_draws(dump, DRAW_IDS)
    if sorted(got) != list(DRAW_IDS):
        raise SystemExit(f"missing draws: {sorted(set(DRAW_IDS) - set(got))}")
    ref = got[DRAW_IDS[0]]["state"]
    for idx in DRAW_IDS:
        st = got[idx]["state"]
        t0, pr = st["tex0"], st["prim"]
        if (t0["tbp"], t0["tbw"], t0["psm"], t0["tw"], t0["th"], t0["tfx"]) != (11264, 4, 0, 256, 256, 0):
            raise SystemExit(f"idx{idx}: unexpected TEX0 {t0}")
        if (pr["type"], pr["tme"], pr["fst"], pr["abe"]) != (6, 1, 0, 1):
            raise SystemExit(f"idx{idx}: unexpected PRIM {pr}")
        for key in ("frame", "test", "zbuf", "alpha", "clamp"):
            if st[key] != ref[key]:
                raise SystemExit(f"idx{idx}: {key} differs")

    mem, *_ = RC.build_localmem_to(dump, DRAW_IDS[0])
    tex = [mem.read_ct32_word(11264, 4, x, y) for y in range(TH) for x in range(TW)]
    crc = sum(tex) & 0xffffffff
    wmem(f"sh3_{tag}0_tex_lpddr.mem", tex,
         f"Ch414 authentic 256x256 PSMCT32 intermediate, linearized; crc=0x{crc:08x}")

    pieces = []
    for idx in DRAW_IDS:
        pieces.extend(sprite_bands(got[idx]))
    epochs, cur, cov = [], [], 0
    for tris, npx in pieces:
        if cur and cov + npx > CAPACITY:
            epochs.append((cur, cov)); cur, cov = [], 0
        cur.extend(tris); cov += npx
    if cur:
        epochs.append((cur, cov))

    tex0 = bake.tex0_pack(NEW_TBP, 4, psm=0, tw=8, th=8, tfx=0)
    prim = 3 | (1 << 4) | (1 << 6)  # independent TRIANGLE records, TME, ABE
    # Preserve authentic RGB-only FBMSK while changing only FBP/FBW.
    frame = bake.frame_1_psmct32(FBW) | (ref["frame"] & 0xffffffff00000000)
    for k, (tris, npx) in enumerate(epochs):
        stg = [len(tris) | (1 << 32) | (1 << 34), frame, ref["alpha"],
               ref["test"], ref["zbuf"], tex0, ref["clamp"], prim]
        for tri in tris:
            for v in tri:
                stg.extend(vertex_words(v))
        if len(stg) > STG_WORDS:
            raise SystemExit(f"epoch{k}: staging overflow {len(stg)}")
        bake.write_feeder_stg_mem(f"feeder_sh3_{tag}{k}.mem", stg,
            f"Ch414 authentic CT32 darken epoch{k}, {len(tris)} tris, coverage {npx}", total=STG_WORDS)

    with open(os.path.join(DATA, f"sh3_{tag}_epochs.txt"), "w") as f:
        f.write("# Ch414 authentic PSMCT32 intermediate darkening sprites\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 {len(epochs)} fbpxw {W} fbh {H} fbwords {W*H} lpddr_tex 0x200000 tex_words {len(tex)} n_beats {N_BEATS}\n")
        for k, (tris, _npx) in enumerate(epochs):
            f.write(f"{k} {DRAW_IDS[0]} 11264 0 sh3_{tag}0_tex_lpddr.mem 0x200000 {TEX_BYTES} 0x{crc:08x} "
                    f"feeder_sh3_{tag}{k}.mem {8 + 9*len(tris)} {len(tris)} {0 if k == 0 else 1} - 0x00000000\n")
    print(f"[Ch414] PASS: {len(DRAW_IDS)} CT32 sprites -> {len(pieces)} bands -> {len(epochs)} epochs; "
          f"max coverage={max(x[1] for x in epochs)} crc=0x{crc:08x}")
    return 0


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