"""Wwise packed codebook library (ww2ogg 'packed_codebooks.bin' format).

The game's built-in table (runtime 0x0158E748, 598 entries) is byte-identical to
ww2ogg's standard packed_codebooks.bin; see docs/wem_format.md.
"""
import os
import struct

from .bits import BitReader, BitWriter, ilog, book_maptype1_quantvals

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_LIB = os.path.join(HERE, "data", "game_codebooks.bin")


def load_library(path: str = DEFAULT_LIB) -> list:
    b = open(path, "rb").read()
    table = struct.unpack_from("<I", b, len(b) - 4)[0]
    n = (len(b) - table) // 4
    offs = struct.unpack_from("<%dI" % n, b, table)
    return [b[offs[i]:offs[i + 1]] for i in range(n - 1)]


def save_library(path: str, books: list):
    out = bytearray()
    offs = []
    for bk in books:
        offs.append(len(out))
        out += bk
    offs.append(len(out))  # table start == end of last book
    out += struct.pack("<%dI" % len(offs), *offs)
    open(path, "wb").write(bytes(out))


def wwise_book_to_vorbis(blob: bytes, w: BitWriter):
    """Expand a compact Wwise codebook into a standard Vorbis codebook (ww2ogg rebuild)."""
    r = BitReader(blob)
    w.put(24, 0x564342)
    dims = r.get(4); w.put(16, dims)
    entries = r.get(14); w.put(24, entries)
    ordered = r.get(1); w.put(1, ordered)
    if ordered:
        w.put(5, r.get(5))
        cur = 0
        while cur < entries:
            nb = ilog(entries - cur)
            num = r.get(nb); w.put(nb, num)
            cur += num
    else:
        cwll = r.get(3)
        sparse = r.get(1); w.put(1, sparse)
        if cwll == 0 or cwll > 5:
            raise ValueError("bad codeword length length")
        for _ in range(entries):
            present = 1
            if sparse:
                present = r.get(1); w.put(1, present)
            if present:
                w.put(5, r.get(cwll))
    lookup = r.get(1); w.put(4, lookup)
    if lookup == 1:
        w.put(32, r.get(32)); w.put(32, r.get(32))
        vlen = r.get(4); w.put(4, vlen)
        w.put(1, r.get(1))
        for _ in range(book_maptype1_quantvals(entries, dims)):
            w.put(vlen + 1, r.get(vlen + 1))
    elif lookup != 0:
        raise ValueError("bad lookup type")
    return r.pos


def skip_vorbis_book(r: BitReader):
    """Advance reader past one standard Vorbis codebook; returns (dims, entries)."""
    if r.get(24) != 0x564342:
        raise ValueError("codebook sync lost")
    dims = r.get(16)
    entries = r.get(24)
    if r.get(1):  # ordered
        r.get(5)
        cur = 0
        while cur < entries:
            cur += r.get(ilog(entries - cur))
    else:
        sparse = r.get(1)
        for _ in range(entries):
            if not sparse or r.get(1):
                r.get(5)
    lookup = r.get(4)
    if lookup in (1, 2):
        r.get(32); r.get(32)
        vlen = r.get(4); r.get(1)
        n = book_maptype1_quantvals(entries, dims) if lookup == 1 else entries * dims
        r.pos += n * (vlen + 1)
    elif lookup != 0:
        raise ValueError("bad lookup type %d" % lookup)
    return dims, entries


def book_bits_key(data: bytes, start: int, end: int) -> tuple:
    """Hashable representation of a bit range [start,end)."""
    n = end - start
    v = int.from_bytes(data, "little") >> start
    return (n, v & ((1 << n) - 1))


class CodebookLibrary:
    def __init__(self, path: str = DEFAULT_LIB):
        self.books = load_library(path)
        self.by_vorbis = {}
        self.vorbis = []
        for i, blob in enumerate(self.books):
            w = BitWriter()
            wwise_book_to_vorbis(blob, w)
            key = (w.pos, int.from_bytes(w.bytes(), "little") & ((1 << w.pos) - 1))
            self.vorbis.append(key)
            self.by_vorbis.setdefault(key, i)

    def lookup(self, key):
        return self.by_vorbis.get(key)