"""Wwise soundbanks (bank version 56): chunks, in-memory media (DIDX/DATA) and HIRC sound sources.

Works on a buffer that may hold a bank anywhere (a loose .bnk file, a BKPK pack's embedded bank, or
an uncompressed .upk package with many banks). Offsets are relative to the start of the bank.
Format: docs/formats/banks.md.
"""
import re
import struct
from collections import namedtuple

CHUNKS = (b"BKHD", b"DIDX", b"DATA", b"HIRC", b"STID", b"STMG", b"ENVS", b"FXPR")
SOUND = 2                         # HIRC object type of a Sound
IN_MEMORY, STREAMED, PREFETCH = 0, 1, 2

Media = namedtuple("Media", "id off size")                     # off: from bank start
Source = namedtuple("Source", "object stream media bank off size pos")  # pos: of `off` field, from bank start


class BankError(ValueError):
    pass


class Bank:
    def __init__(self, buf, at=0):
        self.buf, self.at = buf, at
        self.chunks = {}
        p = at
        while p + 8 <= len(buf) and bytes(buf[p:p + 4]) in CHUNKS:
            tag = bytes(buf[p:p + 4]).decode()
            n = struct.unpack_from("<I", buf, p + 4)[0]
            if p + 8 + n > len(buf) or tag in self.chunks:
                break
            self.chunks[tag] = (p + 8 - at, n)
            p += 8 + n
        if "BKHD" not in self.chunks or self.chunks["BKHD"][1] < 8:
            raise BankError("no bank at 0x%X" % at)
        self.size = p - at
        o = at + self.chunks["BKHD"][0]
        self.version, self.id = struct.unpack_from("<II", buf, o)

    def _abs(self, tag):
        o, n = self.chunks[tag]
        return self.at + o, n

    def media(self):
        """In-memory WEMs: [Media(id, offset from bank start, size)]."""
        if "DIDX" not in self.chunks or "DATA" not in self.chunks:
            return []
        o, n = self._abs("DIDX")
        data = self.chunks["DATA"][0]
        out = []
        for i in range(n // 12):
            mid, off, size = struct.unpack_from("<III", self.buf, o + 12 * i)
            out.append(Media(mid, data + off, size))
        return out

    def sources(self):
        """Sources of the HIRC Sound objects: [Source]. In-memory sources carry offset/size."""
        if "HIRC" not in self.chunks:
            return []
        o, n = self._abs("HIRC")
        end = o + n
        count = struct.unpack_from("<I", self.buf, o)[0]
        p = o + 4
        out = []
        for _ in range(count):
            if p + 5 > end:
                raise BankError("HIRC overruns its chunk in bank %08X" % self.id)
            t = self.buf[p]
            ln = struct.unpack_from("<I", self.buf, p + 1)[0]
            if t == SOUND and ln >= 28:
                obj, _plugin, stream, mid, bank = struct.unpack_from("<5I", self.buf, p + 5)
                off, size = struct.unpack_from("<II", self.buf, p + 25)
                out.append(Source(obj, stream, mid, bank, off, size, p + 25 - self.at))
            p += 5 + ln
        return out

    def names(self):
        """STID: {bank id: name}."""
        if "STID" not in self.chunks:
            return {}
        o, n = self._abs("STID")
        count = struct.unpack_from("<I", self.buf, o + 4)[0]
        p = o + 8
        out = {}
        for _ in range(count):
            bid, ln = struct.unpack_from("<IB", self.buf, p)
            out[bid] = bytes(self.buf[p + 5:p + 5 + ln]).decode("latin-1")
            p += 5 + ln
        return out


def find_banks(buf, start=0):
    """Every bank in `buf` (e.g. an uncompressed package): yields Bank."""
    for m in re.finditer(b"BKHD", buf[start:] if start else buf):
        try:
            yield Bank(buf, m.start() + start)
        except (BankError, struct.error):
            continue


def fnv1(name):
    """Wwise ID of a name (bank, event, ...): FNV-1 32 of the lower-cased name."""
    h = 2166136261
    for c in name.lower().encode("utf-8"):
        h = (h * 16777619) & 0xFFFFFFFF
        h ^= c
    return h