#!/usr/bin/env python3
"""Reference DLX1 reader/decoder -- the ground truth the 68000 player is checked against.

This is deliberately a *decoder*, not a re-run of the encoder: it parses the
container byte-for-byte the way `src/player/` must, so that any disagreement
between it and the 68000 is a decoder bug rather than two encoders differing.
`tools/analysis/09_ratectl_drift.py` already gates the encoder against its own
reconstruction; this gates the container against the player.

Everything is big-endian (see the `encode.py` docstring).  Block raster order,
2-bit modes packed MSB-first: 00=SKIP 01=V1 10=V4 11=RAW.

V4 sub-block order is (sub_y, sub_x) row-major -- TL, TR, BL, BR -- matching
`vq_hybrid.paint`'s reshape(-1,2,2,2,2).transpose(0,1,3,2,4).
"""
import struct
import numpy as np

MODE_SKIP, MODE_V1, MODE_V4, MODE_RAW = 0, 1, 2, 3


class DLX:
    def __init__(self, path):
        self.raw = open(path, "rb").read()
        b = self.raw
        if b[:4] != b"DLX1":
            raise ValueError(f"{path}: not a DLX1 container")
        (self.W, self.H, self.fps, self.nframes,
         self.k1, self.k4) = struct.unpack(">HHHHHH", b[4:16])
        off_pal, off_cb1, off_cb4, off_frm = struct.unpack(">IIII", b[16:32])

        self.pal = np.frombuffer(b, np.uint8, 256 * 3, off_pal).reshape(256, 3)
        self.cb1 = np.frombuffer(b, np.uint8, self.k1 * 16,
                                 off_cb1).reshape(self.k1, 4, 4)
        self.cb4 = np.frombuffer(b, np.uint8, self.k4 * 4,
                                 off_cb4).reshape(self.k4, 2, 2)

        self.idx_bytes = 1 if max(self.k1, self.k4) <= 256 else 2
        self.nbx, self.nby = self.W // 4, self.H // 4
        self.nb = self.nbx * self.nby
        self.mode_bytes = (self.nb * 2 + 7) // 8

        # frame directory: (offset of the mode header, payload length)
        self.frames = []
        p = off_frm
        for _ in range(self.nframes):
            (n,) = struct.unpack(">I", b[p:p + 4])
            self.frames.append((p + 4, n))
            p += 4 + n
        if p != len(b):
            raise ValueError(f"{path}: {len(b) - p} trailing bytes after "
                             f"{self.nframes} frames")

    def modes(self, f):
        o, _ = self.frames[f]
        h = np.frombuffer(self.raw, np.uint8, self.mode_bytes, o)
        m = np.stack([(h >> 6) & 3, (h >> 4) & 3, (h >> 2) & 3, h & 3], axis=1)
        return m.reshape(-1)[:self.nb].copy()

    def blocks(self, f):
        """Decoded 4x4 palette-index blocks for the non-SKIP blocks of frame f.

        Returns (mode, dict{block index -> (4,4) uint8}).  SKIP blocks are
        absent by construction -- the player must leave those pixels alone,
        and a decoder that materialises them is hiding the very bug this
        module exists to catch.
        """
        mode = self.modes(f)
        o, n = self.frames[f]
        p, end = o + self.mode_bytes, o + n
        ib, out = self.idx_bytes, {}
        b = self.raw
        for i, mo in enumerate(mode):
            if mo == MODE_SKIP:
                continue
            if mo == MODE_V1:
                v = b[p] if ib == 1 else (b[p] << 8) | b[p + 1]
                p += ib
                out[i] = self.cb1[v]
            elif mo == MODE_V4:
                sub = []
                for _ in range(4):
                    v = b[p] if ib == 1 else (b[p] << 8) | b[p + 1]
                    p += ib
                    sub.append(self.cb4[v])
                blk = np.empty((4, 4), np.uint8)
                blk[0:2, 0:2], blk[0:2, 2:4] = sub[0], sub[1]
                blk[2:4, 0:2], blk[2:4, 2:4] = sub[2], sub[3]
                out[i] = blk
            else:
                out[i] = np.frombuffer(b, np.uint8, 16, p).reshape(4, 4)
                p += 16
        if p != end:
            raise ValueError(f"frame {f}: payload consumed {p - o} of {n} bytes")
        return mode, out

    def paint(self, canvas, f):
        """Apply frame f in place to a (H,W) index canvas.  SKIP = untouched."""
        mode, blks = self.blocks(f)
        for i, blk in blks.items():
            by, bx = divmod(i, self.nbx)
            canvas[by * 4:by * 4 + 4, bx * 4:bx * 4 + 4] = blk
        return mode

    def decode_all(self):
        """The true reconstruction sequence: what any correct player displays."""
        c = np.zeros((self.H, self.W), np.uint8)
        out = []
        for f in range(self.nframes):
            self.paint(c, f)
            out.append(c.copy())
        return out