"""Wwise RIFF (.wem) container for Arkham City (Wwise 2011.1, fmt size 0x42, Vorbis 0xFFFF)."""
import re
import struct
from dataclasses import dataclass, field

VORB_FMT = "<IIIHHIIHHIII2B"  # AkVorbisInfo as stored at fmt+0x18 (42 bytes)
VORB_FIELDS = ["total_frames", "loop_start_pkt", "loop_end_pkt", "loop_begin_extra", "loop_end_extra",
               "seek_table_size", "vorbis_data_offset", "max_packet_size", "last_granule_extra",
               "decode_alloc", "decode_alloc_x64", "codebook_hash", "bs0_exp", "bs1_exp"]


@dataclass
class Chunk:
    tag: bytes
    data: bytes


@dataclass
class Wem:
    chunks: list
    channels: int = 0
    rate: int = 0
    avg_bps: int = 0
    valid_bits: int = 0
    channel_mask: int = 0
    vorb: dict = field(default_factory=dict)

    # ---- parsing ----
    @classmethod
    def parse(cls, b: bytes) -> "Wem":
        if b[:4] != b"RIFF" or b[8:12] != b"WAVE":
            raise ValueError("not RIFF/WAVE")
        riff_end = 8 + struct.unpack_from("<I", b, 4)[0]
        pos = 12
        chunks = []
        while pos + 8 <= min(riff_end, len(b)):
            tag, ln = struct.unpack_from("<4sI", b, pos)
            chunks.append(Chunk(tag, b[pos + 8:pos + 8 + ln]))
            pos += 8 + ln + (ln & 1)
        w = cls(chunks)
        fmt = w.get(b"fmt ")
        if fmt is None or len(fmt) != 0x42:
            raise ValueError("not a Wwise Vorbis WEM (fmt size %s)" % (None if fmt is None else hex(len(fmt))))
        fmt_tag, w.channels, w.rate, w.avg_bps, balign, bps, cb = struct.unpack_from("<HHIIHHH", fmt, 0)
        if fmt_tag != 0xFFFF or len(fmt) != 0x42 or cb != 0x30:
            raise ValueError("unsupported fmt (tag=%x size=%x)" % (fmt_tag, len(fmt)))
        w.valid_bits, w.channel_mask = struct.unpack_from("<HI", fmt, 0x12)
        w.vorb = dict(zip(VORB_FIELDS, struct.unpack_from(VORB_FMT, fmt, 0x18)))
        return w

    def get(self, tag: bytes):
        for c in self.chunks:
            if c.tag == tag:
                return c.data
        return None

    def set(self, tag: bytes, data: bytes):
        for c in self.chunks:
            if c.tag == tag:
                c.data = data
                return
        self.chunks.append(Chunk(tag, data))

    # ---- data chunk layout ----
    @property
    def data(self) -> bytes:
        return self.get(b"data")

    def setup_packet(self) -> bytes:
        d = self.data
        off = self.vorb["seek_table_size"]
        ln = struct.unpack_from("<H", d, off)[0]
        return d[off + 2: off + 2 + ln]

    def audio_packets(self):
        """List of (offset_in_data, payload) for audio packets (2-byte size headers)."""
        d = self.data
        pos = self.vorb["vorbis_data_offset"]
        out = []
        while pos + 2 <= len(d):
            ln = struct.unpack_from("<H", d, pos)[0]
            if ln == 0 or pos + 2 + ln > len(d):
                break
            out.append((pos, d[pos + 2: pos + 2 + ln]))
            pos += 2 + ln
        return out, pos

    # ---- serialization ----
    def build_fmt(self) -> bytes:
        v = self.vorb
        return (struct.pack("<HHIIHHH", 0xFFFF, self.channels, self.rate, self.avg_bps, 0, 0, 0x30)
                + struct.pack("<HI", self.valid_bits, self.channel_mask)
                + struct.pack(VORB_FMT, *[v[k] for k in VORB_FIELDS]))

    def to_bytes(self) -> bytes:
        self.set(b"fmt ", self.build_fmt())
        body = b"WAVE"
        for c in self.chunks:
            # the game's WEMs never pad odd-sized chunks (all chunks before 'data' are even)
            body += struct.pack("<4sI", c.tag, len(c.data)) + c.data
        return b"RIFF" + struct.pack("<I", len(body)) + body


# ---- filling an in-memory slot exactly ----
# The game ships no in-memory WEM with bytes after its RIFF inside its slot, and plays a replacement that is
# followed by zero padding there as silence (findings 24). A smaller replacement is therefore grown to the
# slot's exact size with a label list chunk before `data`, the form the game's own WEMs use for markers:
# LIST{'adtl', labl{u32 cue id 0, zero text}}. Even sizes only (the game's chunks before `data` are even).
FILL_MIN = 8 + 4 + 8 + 4 + 2


def fillable(length, size):
    """True if a WEM of `length` bytes can fill a slot of `size` bytes exactly."""
    pad = size - length
    return pad == 0 or (pad >= FILL_MIN and pad % 2 == 0)


def fill(wem: bytes, size: int) -> bytes:
    """`wem` grown to exactly `size` bytes (see above). ValueError if `fillable` is False."""
    if not fillable(len(wem), size):
        raise ValueError("a %d-byte WEM cannot fill a %d-byte slot exactly" % (len(wem), size))
    pad = size - len(wem)
    if not pad:
        return wem
    i = _data_chunk(wem)
    text = pad - 8 - 4 - 8 - 4
    chunk = (b"LIST" + struct.pack("<I", pad - 8) + b"adtl" + b"labl" + struct.pack("<II", text + 4, 0)
             + bytes(text))
    out = wem[:i] + chunk + wem[i:]
    return out[:4] + struct.pack("<I", len(out) - 8) + out[8:]


def unfill(wem: bytes) -> bytes:
    """The WEM without a filler chunk added by `fill` (a game WEM is returned unchanged)."""
    i = _data_chunk(wem)
    if i is None:
        return wem
    for j in range(12, i, 2):
        if wem[j:j + 4] == b"LIST" and wem[j + 8:j + 16] == b"adtllabl":
            n = 8 + struct.unpack_from("<I", wem, j + 4)[0]
            if j + n == i and struct.unpack_from("<I", wem, j + 20)[0] == 0 and not wem[j + 24:i].strip(b"\0"):
                out = wem[:j] + wem[i:]
                return out[:4] + struct.pack("<I", len(out) - 8) + out[8:]
    return wem


def _data_chunk(wem: bytes):
    """Offset of the `data` chunk header (walking the chunks), or None."""
    pos = 12
    while pos + 8 <= len(wem):
        tag, ln = struct.unpack_from("<4sI", wem, pos)
        if tag == b"data":
            return pos
        pos += 8 + ln
    return None


SUB_ALIGN = 0x800
_NONZERO = re.compile(rb"[^\x00]")


def split_concatenated(b: bytes):
    """All WEMs of a WAD entry, as [(offset, riff_bytes)], plus the end of the last one.

    LOC entries group variations of a line (e.g. CHR_RadioScanner = 63 WEMs): each WEM starts at the
    next 0x800 boundary after the previous one (zero padding); a few entries pack them back to back."""
    out = []
    pos = 0
    end = 0
    while pos + 12 <= len(b) and b[pos:pos + 4] == b"RIFF" and b[pos + 8:pos + 12] == b"WAVE":
        end = pos + 8 + struct.unpack_from("<I", b, pos + 4)[0]
        out.append((pos, b[pos:end]))
        nxt = (end + SUB_ALIGN - 1) // SUB_ALIGN * SUB_ALIGN
        if b[end:end + 4] == b"RIFF":
            pos = end
        elif b[nxt:nxt + 4] == b"RIFF":
            pos = nxt
        else:  # zero padding of arbitrary length (e.g. after an in-place replacement)
            m = _NONZERO.search(b, end)
            if not m:
                break
            pos = m.start()
    return out, end