"""Model of the game's codec-setup bump-allocator usage (dwDecodeAllocSize).

Reversed from the dumped BatmanAC.exe (see docs/reverse-engineering/decoder.md):
WwiseVorbisParseSetupHeader (runtime 0x010B4D10) and its callees are a Tremor-lowmem
derivative; every pool allocation is rounded up to 4 bytes.
"""
from .bits import BitReader, ilog, book_maptype1_quantvals
from .codebooks import CodebookLibrary


def a4(n):
    return (n + 3) & ~3


def _node_bytes(used, leafwidth):
    if used < 2:
        return 4
    if leafwidth == 3:
        leafwidth = 4
    if ilog(3 * used - 6) + 1 <= leafwidth * 4:
        return leafwidth // 2 if leafwidth // 2 else 1
    return leafwidth


def _leaf_words(nodeb, leafwidth):
    return 2 if leafwidth > nodeb else 1


def _decode_table(used, nodeb, leafw):
    if nodeb == 4:
        return used * 8 + 4
    return a4((used * (leafw + 1) - 2) * nodeb)


def codebook_alloc(blob: bytes) -> int:
    """Pool bytes used when unpacking one compact Wwise codebook."""
    r = BitReader(blob)
    dims = r.get(4)
    entries = r.get(14)
    if r.get(1):  # ordered
        r.get(5)
        cur = 0
        while cur < entries:
            cur += r.get(ilog(entries - cur))
        used = entries
    else:
        cwll = r.get(3)
        sparse = r.get(1)
        used = 0
        for _ in range(entries):
            if not sparse or r.get(1):
                r.get(cwll)
                used += 1
    maptype = r.get(1)
    if maptype == 0:
        lw = ilog(entries) // 8 + 1
        nodeb = _node_bytes(used, lw)
        return _decode_table(used, nodeb, _leaf_words(nodeb, lw))
    r.get(32); r.get(32)
    qbits = r.get(4) + 1
    quantvals = book_maptype1_quantvals(entries, dims)
    total1 = (qbits * dims + 8) // 8
    total2 = (ilog(quantvals - 1) * dims + 8) // 8 + (qbits + 7) // 8
    if total1 <= 4 and total1 <= total2:
        nodeb = _node_bytes(used, total1)
        return _decode_table(used, nodeb, _leaf_words(nodeb, total1))
    qv = a4(quantvals) if qbits <= 8 else a4(quantvals * 2)
    lw = (ilog(quantvals - 1) * dims + 8) // 8
    nodeb = _node_bytes(used, lw)
    return qv + _decode_table(used, nodeb, _leaf_words(nodeb, lw))


def setup_alloc(ww_setup: bytes, channels: int, lib: CodebookLibrary) -> dict:
    """Walk a stripped Wwise setup and sum pool allocations, by section."""
    r = BitReader(ww_setup)
    out = {}
    nb = r.get(8) + 1
    total = nb * 0x3C
    cb = 0
    for _ in range(nb):
        cb += codebook_alloc(lib.books[r.get(10)])
    out["codebooks"] = total + cb
    nf = r.get(6) + 1
    fl = nf * 0x24
    for _ in range(nf):
        parts = r.get(5)
        classes = [r.get(4) for _ in range(parts)]
        maxc = max(classes) if classes else -1
        dims = []
        for _ in range(maxc + 1):
            dims.append(r.get(3) + 1)
            sub = r.get(2)
            if sub:
                r.get(8)
            for _ in range(1 << sub):
                r.get(8)
        r.get(2)
        rb = r.get(4)
        posts = sum(dims[c] for c in classes)
        r.pos += posts * rb
        fl += a4(parts) + a4(11 * (maxc + 1)) + a4(2 * (posts + 2)) + a4(posts + 2) + 2 * a4(posts)
    out["floors"] = fl
    floors_count = nf
    nr = r.get(6) + 1
    rs = nr * 0x1C
    for _ in range(nr):
        r.get(2); r.get(24); r.get(24); r.get(24)
        ncls = r.get(6) + 1
        r.get(8)
        casc = []
        for _ in range(ncls):
            lo = r.get(3)
            hi = r.get(5) if r.get(1) else 0
            casc.append(hi * 8 + lo)
        for c in casc:
            r.pos += 8 * bin(c & 0xFF).count("1")
        rs += a4(ncls) + 8 * ncls
    out["residues"] = rs
    nm = r.get(6) + 1
    ms = nm * 0x14
    for _ in range(nm):
        submaps = r.get(4) + 1 if r.get(1) else 1
        if r.get(1):
            steps = r.get(8) + 1
            b = ilog(channels - 1)
            r.pos += steps * 2 * b
            ms += a4(2 * steps)
        r.get(2)
        if submaps > 1:
            r.pos += 4 * channels
            ms += a4(channels)
        r.pos += 24 * submaps
        ms += a4(2 * submaps)
    out["mappings"] = ms
    nmodes = r.get(6) + 1
    out["modes"] = a4(2 * nmodes)
    out["total"] = sum(out.values())
    out["floors_count"] = floors_count
    return out


def fnv1_32(b: bytes) -> int:
    """uHashCodebook: FNV-1 32-bit over the stripped setup packet (verified on all 24 game setups)."""
    h = 0x811C9DC5
    for x in b:
        h = (h * 0x01000193) & 0xFFFFFFFF
        h ^= x
    return h


def decode_alloc(ww_setup: bytes, channels: int, lib: CodebookLibrary):
    """(dwDecodeAllocSize, dwDecodeX64AllocSize) as the 2011.1 encoder writes them.

    32-bit: 0x30 (setup info) + parser allocations + 4*nbooks + 8 -> exact for all 24 game setups.
    x64: fitted (alloc + 8*nbooks + 40 + 52*nfloors); not read by the 32-bit game.
    """
    a = setup_alloc(ww_setup, channels, lib)
    nbooks = ww_setup[0] + 1
    alloc = 0x30 + a["total"] + 4 * nbooks + 8
    nfloors = a["floors_count"]
    return alloc, alloc + 8 * nbooks + 40 + 52 * nfloors