"""BKPK stream packs: SFX\\<ID>.bnk ("BKPK" header + embedded bank) indexes the streamed WEMs
stored inside the SFX WAD entry with the same ID (e.g. EDF036D6 'Music', E8F980BE 'Music_PDLC',
57572DAC 'Global', FCA4F778).

BKPK header (u32 LE):
  0x00 'BKPK'   0x04 version 0x01020000   0x08 pack id (== WAD entry hash)
  0x10 media count N   0x18 table size (12N)   0x1C table offset (0x48)
  0x20 total stream size (== WAD entry size; end of last media aligned to 0x800)
  0x2C end of table (0x48 + 12N)   0x38 embedded bank size   0x3C embedded bank offset
  table: N x {u32 media id, u32 size, u32 offset-in-WAD-entry}, sorted by id; each media
  starts at the next 0x800-aligned offset after the previous one.
See docs/formats/wad.md.
"""
import struct

ALIGN = 0x800


def _a(x):
    return (x + ALIGN - 1) // ALIGN * ALIGN


class Pack:
    def __init__(self, bnk: bytes):
        if bnk[:4] != b"BKPK":
            raise ValueError("not a BKPK pack")
        self.raw = bnk
        self.id = struct.unpack_from("<I", bnk, 8)[0]
        n = struct.unpack_from("<I", bnk, 0x10)[0]
        self.table_off = struct.unpack_from("<I", bnk, 0x1C)[0]
        self.total = struct.unpack_from("<I", bnk, 0x20)[0]
        self.entries = [list(struct.unpack_from("<III", bnk, self.table_off + 12 * i)) for i in range(n)]
        self.index = {e[0]: i for i, e in enumerate(self.entries)}

    def media(self, payload: bytes, mid: int) -> bytes:
        _, size, off = self.entries[self.index[mid]]
        return payload[off:off + size]

    def gaps(self, payload: bytes) -> dict:
        """Non-zero padding between media: {start offset: bytes}. Saved with the originals so that a
        full revert restores the pack byte-for-byte."""
        out = {}
        for i, (mid, size, off) in enumerate(self.entries):
            a = off + size
            b = self.entries[i + 1][2] if i + 1 < len(self.entries) else self.total
            g = payload[a:b]
            if g.strip(b"\0"):
                out[a] = g
        return out

    def rebuild(self, payload: bytes, replacements: dict, gaps: dict = None):
        """Returns (new_bnk_bytes, new_wad_payload).

        Media are laid out by the pack rules (sorted order, 0x800-aligned, total aligned). Padding is
        taken from `gaps` (original non-zero padding) when its position matches, else from the current
        payload when that gap did not move, else zeros."""
        missing = set(replacements) - set(self.index)
        if missing:
            raise KeyError("media not in pack: %s" % ", ".join("%08X" % m for m in missing))
        gaps = gaps or {}
        cur_gaps = {}
        for i, (mid, size, off) in enumerate(self.entries):
            nxt = self.entries[i + 1][2] if i + 1 < len(self.entries) else self.total
            cur_gaps[off + size] = nxt

        out = bytearray()
        new_entries = []

        def pad(upto):
            a, n = len(out), upto - len(out)
            if n <= 0:
                return
            g = gaps.get(a)
            if g is not None and len(g) == n:
                out.extend(g)
            elif cur_gaps.get(a) == upto:
                out.extend(payload[a:upto])
            else:
                out.extend(bytes(n))

        for mid, size, off in self.entries:
            pad(_a(len(out)))
            blob = replacements[mid] if mid in replacements else payload[off:off + size]
            new_entries.append((mid, len(blob), len(out)))
            out += blob
        total = _a(len(out))
        pad(total)
        bnk = bytearray(self.raw)
        struct.pack_into("<I", bnk, 0x20, total)
        for i, e in enumerate(new_entries):
            struct.pack_into("<III", bnk, self.table_off + 12 * i, *e)
        return bytes(bnk), bytes(out)