#!/usr/bin/env python3
"""retroDE_ps2 — Ch349 steps 2-4: reconstruct ONE real SH3 textured draw from GS local memory.

Closes the Ch347/348 gap (authentic asset on CHOSEN geometry) -> an ACTUAL commercial draw reconstructed
faithfully. Default pick (gs_sh3_draw_census.py top): frame 1 idx89761, a 70-prim PSMT8 512x512 TRI_STRIP at
tbp=9216, CLUT cbp=13952 (CSM1 PSMCT32). That texture is STREAMED as 256x256 PSMCT32 (upload idx13288, same
262144 bytes) and SAMPLED as 512x512 PSMT8 — the exact stream-one-format / sample-another bridge.

  step 2  Build a GS local-memory model (gs_localmem.LocalMem) seeded from the dump's initial VRAM snapshot,
          replay every host->local PSMCT32 upload up to the draw, then READ the texture back via the PSMT8
          swizzle. index -> CLUT -> ABGR. This is "the texture as the real draw sees it".
  step 3  decode + print the draw's real TEX0/CLUT/state.
  step 4  (gs_sh3_draw_ref.py) rasterize the actual geometry sampling this texture.

All outputs are SH3-derived -> LOCAL/gitignored (captures/gs/silenthill3/extracted/recon/).
Usage:  gs_sh3_recon.py <dump.gs.zst> [--draw-idx N] [--tbp T] [--cbp C] [--tbw W] [--tw 512] [--th 512]
                                       [--clut-order linear|grid] [--out DIR]
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gs_parse
import gs_texture_residency as R
import gs_localmem as LM

def build_localmem_to(dump, draw_idx):
    """Seed VRAM from the initial snapshot and replay all host->local PSMCT32 uploads with idx < draw_idx.
    Returns (mem, replayed_uploads, all_uploads, events, vram)."""
    d, h, events, uploads, runs, vram = R.collect(dump, 0)
    if vram is None:
        return None, [], uploads, events, None
    mem = LM.LocalMem(vram)
    replayed = []
    for u in uploads:
        if u["idx"] >= draw_idx: continue
        if u["dpsm"] != 0x00:        # only PSMCT32 stream writes modelled here (SH3 env path)
            continue
        off, end = u["blob_range"]
        blob = d[off:end]
        words = [int.from_bytes(blob[i*4:i*4+4], "little") for i in range(len(blob)//4)]
        mem.write_image_ct32(u["dbp"], u["dbw"], u["dx"], u["dy"], u["w"], u["h"], words)
        replayed.append(u)
    return mem, replayed, uploads, events, vram

def read_clut32(mem, cbp, order="grid"):
    """Read a 256-entry PSMCT32 CLUT from the modelled VRAM. 'grid' = read as a 16x16 CT32 surface based at
    cbp (dbw=1) — the layout a CSM1 8-bit palette occupies; 'linear' = raw contiguous i*4 from cbp*256.
    Returns 256 packed ints (PS2 PSMCT32 word == 0xAABBGGRR, low byte R)."""
    pal = [0]*256
    if order == "linear":
        base = cbp*256
        for i in range(256):
            a = base + i*4
            pal[i] = int.from_bytes(mem.m[a:a+4], "little") if a+4 <= mem.SIZE else 0
    else:  # grid: palette entry i at (x=i%16, y=i//16) via CT32 swizzle, dbw=1
        for i in range(256):
            pal[i] = mem.read_ct32_word(cbp, 1, i & 15, i >> 4)
    return pal

def decode_pixel(pal, idx):
    p = pal[idx & 0xFF]
    return (p & 0xFF, (p>>8)&0xFF, (p>>16)&0xFF, (p>>24)&0xFF)   # R,G,B,A

def save_png(path, w, h, rgba_pixels):
    from PIL import Image
    img = Image.new("RGBA", (w, h)); img.putdata(rgba_pixels); img.save(path)

def main(argv):
    if len(argv) < 2:
        print(__doc__); return 2
    dump = argv[1]
    def opt(n, dv=None): return argv[argv.index(n)+1] if n in argv else dv
    draw_idx = int(opt("--draw-idx","89761"))
    tbp = int(opt("--tbp","9216")); cbp = int(opt("--cbp","13952"))
    fbw = int(opt("--tbw","8")); tw = int(opt("--tw","512")); th = int(opt("--th","512"))
    order = opt("--clut-order","grid")
    outdir = opt("--out", os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                                       "captures","gs","silenthill3","extracted","recon"))
    os.makedirs(outdir, exist_ok=True)

    mem, replayed, uploads, events, vram = build_localmem_to(dump, draw_idx)
    if mem is None:
        print("VRAM snapshot ABSENT — cannot reconstruct"); return 1
    print(f"[Ch349] dump={os.path.basename(dump)}  draw_idx={draw_idx}  tbp={tbp} cbp={cbp} fbw={fbw} {tw}x{th}")
    print(f"[step2] GS local-mem seeded from initial snapshot + replayed {len(replayed)} PSMCT32 upload(s) "
          f"before the draw:")
    for u in replayed:
        print(f"          idx{u['idx']} dbp={u['dbp']} dbw={u['dbw']} {u['w']}x{u['h']} {u['bytes']}B")

    idx = mem.read_psmt8(tbp, fbw, tw, th)
    distinct = len(set(idx))
    print(f"[step2] de-swizzled PSMT8 index image: {tw}x{th}, {distinct} distinct indices")
    save_png(os.path.join(outdir, "recon_indices_gray.png"), tw, th, [(b,b,b,255) for b in idx])

    pal = read_clut32(mem, cbp, order=order)
    print(f"[step2] CLUT @cbp={cbp} order={order}: {len(set(pal))} distinct ABGR entries")
    color = [decode_pixel(pal, b) for b in idx]
    save_png(os.path.join(outdir, f"recon_texture_{order}.png"), tw, th, [(r,g,b,255) for (r,g,b,a) in color])
    save_png(os.path.join(outdir, f"recon_clut_{order}.png"), 16, 16,
             [(r,g,b,255) for (r,g,b,a) in (decode_pixel(pal,i) for i in range(256))])
    print(f"[step2] wrote recon_indices_gray.png, recon_texture_{order}.png, recon_clut_{order}.png -> {outdir}")
    return 0

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