#!/usr/bin/env python3
"""retroDE_ps2 — Ch349 GS local-memory model (the bridge Codex framed).

A faithful byte-level model of GS local memory (4 MiB VRAM) so a draw's SAMPLED texture can be reconstructed
as the real hardware sees it — even when the asset is STREAMED in one PSM (e.g. PSMCT32, fast word writes)
and SAMPLED in another (e.g. PSMT8 indexed). Both operations address the SAME physical bytes through their
respective GS swizzles; modeling that crossover is the whole point.

Swizzle math is ported verbatim from the project's own RTL (which is in turn locked to PCSX2 GSTables.cpp):
  - PSMCT32 write/read  <- rtl/gif_gs/gs_swizzle_psmct32_stub.sv  (block grid + byte_in_block = yb*32+xb*4)
  - PSMT8   read        <- rtl/gif_gs/gs_swizzle_psmt8_stub.sv     (block grid + 16x16 columnTable8)
Address convention matches the codebase: VRAM byte base = PTR*256 (TBP0/DBP/CBP are 14-bit, *256 == 4 MiB),
and PTRs are page-aligned (multiple of 32) so page_index*8192 composes correctly off that base.
"""

import os
import re

# block grid is shared by PSMCT32 and PSMT8 (4 rows x 8 cols), value = block index within page
BLOCK = [
    [ 0, 1, 4, 5,16,17,20,21],
    [ 2, 3, 6, 7,18,19,22,23],
    [ 8, 9,12,13,24,25,28,29],
    [10,11,14,15,26,27,30,31],
]
# PSMT8 within-block 16x16 -> byte permutation (columnTable8)
COL8 = [
    [  0,  4, 16, 20, 32, 36, 48, 52,  2,  6, 18, 22, 34, 38, 50, 54],
    [  8, 12, 24, 28, 40, 44, 56, 60, 10, 14, 26, 30, 42, 46, 58, 62],
    [ 33, 37, 49, 53,  1,  5, 17, 21, 35, 39, 51, 55,  3,  7, 19, 23],
    [ 41, 45, 57, 61,  9, 13, 25, 29, 43, 47, 59, 63, 11, 15, 27, 31],
    [ 96,100,112,116, 64, 68, 80, 84, 98,102,114,118, 66, 70, 82, 86],
    [104,108,120,124, 72, 76, 88, 92,106,110,122,126, 74, 78, 90, 94],
    [ 65, 69, 81, 85, 97,101,113,117, 67, 71, 83, 87, 99,103,115,119],
    [ 73, 77, 89, 93,105,109,121,125, 75, 79, 91, 95,107,111,123,127],
    [128,132,144,148,160,164,176,180,130,134,146,150,162,166,178,182],
    [136,140,152,156,168,172,184,188,138,142,154,158,170,174,186,190],
    [161,165,177,181,129,133,145,149,163,167,179,183,131,135,147,151],
    [169,173,185,189,137,141,153,157,171,175,187,191,139,143,155,159],
    [224,228,240,244,192,196,208,212,226,230,242,246,194,198,210,214],
    [232,236,248,252,200,204,216,220,234,238,250,254,202,206,218,222],
    [193,197,209,213,225,229,241,245,195,199,211,215,227,231,243,247],
    [201,205,217,221,233,237,249,253,203,207,219,223,235,239,251,255],
]

def ct32_addr(dbp, dbw, x, y):
    """Byte address of PSMCT32 pixel (x,y) in a buffer based at dbp (256-byte units), width dbw (64px units)."""
    page_index = (y >> 5) * dbw + (x >> 6)
    block_idx = BLOCK[(y >> 3) & 3][(x >> 3) & 7]
    return dbp*256 + page_index*8192 + block_idx*256 + (y & 7)*32 + (x & 7)*4

def psmt8_addr(tbp, fbw, x, y):
    """Byte address of PSMT8 pixel (x,y) in a buffer based at tbp (256-byte units), width fbw (64px units)."""
    page_index = (y >> 6) * (fbw >> 1) + (x >> 7)
    block_idx = BLOCK[(y >> 4) & 3][(x >> 4) & 7]
    return tbp*256 + page_index*8192 + block_idx*256 + COL8[y & 15][x & 15]

_COL4 = None
_BLOCK4 = (
    (0,2,8,10), (1,3,9,11), (4,6,12,14), (5,7,13,15),
    (16,18,24,26), (17,19,25,27), (20,22,28,30), (21,23,29,31),
)

def _col4_table():
    """Load the canonical PSMT4 permutation from the RTL contract."""
    global _COL4
    if _COL4 is not None:
        return _COL4
    rtl = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                       "rtl", "gif_gs", "gs_swizzle_psmt4_stub.sv")
    text = open(rtl, encoding="utf-8").read()
    lo = text.index("function automatic logic [8:0] col_idx_psmt4")
    hi = text.index("endfunction", lo)
    tab = [None] * 512
    for key, value in re.findall(r"9'd(\d+)\s*:\s*return\s+9'd(\d+)", text[lo:hi]):
        tab[int(key)] = int(value)
    tab[511] = 511  # explicit default arm in the RTL function
    if any(v is None for v in tab) or sorted(tab) != list(range(512)):
        raise RuntimeError("PSMT4 RTL column table is incomplete or non-bijective")
    _COL4 = tuple(tab)
    return _COL4

def psmt4_addr(tbp, fbw, x, y):
    """Return (byte address, high-nibble select) for a PSMT4 texel."""
    if fbw & 1:
        raise ValueError(f"PSMT4 TBW must be even, got {fbw}")
    page_index = (y >> 7) * (fbw >> 1) + (x >> 7)
    block_idx = _BLOCK4[(y >> 4) & 7][(x >> 5) & 3]
    nib = _col4_table()[((y & 15) << 5) | (x & 31)]
    return tbp*256 + page_index*8192 + block_idx*256 + (nib >> 1), bool(nib & 1)

class LocalMem:
    """4 MiB GS VRAM. Seed from the dump's initial snapshot, then replay host->local uploads in order."""
    SIZE = 0x400000
    def __init__(self, init_bytes=None):
        if init_bytes is not None and len(init_bytes) >= self.SIZE:
            self.m = bytearray(init_bytes[:self.SIZE])
        else:
            self.m = bytearray(self.SIZE)

    def write_image_ct32(self, dbp, dbw, dsax, dsay, w, h, words):
        """Host->local upload in PSMCT32: raster-order words fill (dsax..+w)x(dsay..+h) via ct32 swizzle.
        `words` may be shorter than w*h (partial transfer); fill stops when exhausted (GS behaviour)."""
        n = len(words); i = 0
        for py in range(h):
            for px in range(w):
                if i >= n: return
                a = ct32_addr(dbp, dbw, dsax+px, dsay+py)
                if 0 <= a and a+4 <= self.SIZE:
                    self.m[a:a+4] = (words[i] & 0xFFFFFFFF).to_bytes(4, "little")
                i += 1

    def read_psmt8(self, tbp, fbw, tw, th):
        out = bytearray(tw*th)
        for y in range(th):
            r = y*tw
            for x in range(tw):
                a = psmt8_addr(tbp, fbw, x, y)
                out[r+x] = self.m[a] if 0 <= a < self.SIZE else 0
        return out

    def read_psmt4(self, tbp, fbw, tw, th):
        """Return one unpacked 0..15 index byte per raster-order texel."""
        out = bytearray(tw*th)
        for y in range(th):
            r = y*tw
            for x in range(tw):
                a, hi = psmt4_addr(tbp, fbw, x, y)
                b = self.m[a] if 0 <= a < self.SIZE else 0
                out[r+x] = ((b >> 4) & 0xF) if hi else (b & 0xF)
        return out

    def read_ct32_word(self, dbp, dbw, x, y):
        a = ct32_addr(dbp, dbw, x, y)
        return int.from_bytes(self.m[a:a+4], "little") if a+4 <= self.SIZE else 0
