"""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


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