"""Wwise 2011 ("v56"-style) Vorbis <-> standard Vorbis conversion.

Wwise stores a stripped setup header (codebooks referenced by 10-bit IDs into the
built-in library, narrowed fields) and "modified" audio packets (no packet-type bit,
no prev/next window flags). See docs/wem_format.md.
"""
from .bits import BitReader, BitWriter, ilog
from .codebooks import CodebookLibrary, wwise_book_to_vorbis, skip_vorbis_book, book_bits_key

W2V = "w2v"  # Wwise -> Vorbis
V2W = "v2w"  # Vorbis -> Wwise


class SetupInfo:
    def __init__(self):
        self.codebook_ids = []
        self.codebook_dims_entries = []
        self.mode_blockflags = []
        self.floor_count = 0
        self.residue_count = 0
        self.mapping_count = 0

    @property
    def mode_bits(self):
        return ilog(len(self.mode_blockflags) - 1)


def transcode_setup(src: bytes, direction: str, channels: int, lib: CodebookLibrary) -> (bytes, SetupInfo):
    """Convert a setup packet between the Wwise stripped form and the standard Vorbis form."""
    info = SetupInfo()
    w = BitWriter()
    if direction == W2V:
        r = BitReader(src)
        w.put(8, 5)
        for c in b"vorbis":
            w.put(8, c)
    else:
        if src[:7] != b"\x05vorbis":
            raise ValueError("not a vorbis setup packet")
        r = BitReader(src, 7 * 8)

    def f(ww_bits, vb_bits=None):
        """Transcode one field; returns its value."""
        if vb_bits is None:
            vb_bits = ww_bits
        if direction == W2V:
            v = r.get(ww_bits); w.put(vb_bits, v)
        else:
            v = r.get(vb_bits); w.put(ww_bits, v)
            if v >> ww_bits:
                raise ValueError("value %d does not fit in %d Wwise bits" % (v, ww_bits))
        return v

    def const(vb_bits, value, what):
        """Field present only in Vorbis form, with a fixed value."""
        if direction == W2V:
            w.put(vb_bits, value)
        else:
            v = r.get(vb_bits)
            if v != value:
                raise ValueError("%s must be %d for Wwise, got %d" % (what, value, v))

    # codebooks
    nbooks = f(8) + 1
    for _ in range(nbooks):
        if direction == W2V:
            cid = r.get(10)
            if cid >= len(lib.books):
                raise ValueError("codebook id %d out of range" % cid)
            wwise_book_to_vorbis(lib.books[cid], w)
        else:
            start = r.pos
            de = skip_vorbis_book(r)
            info.codebook_dims_entries.append(de)
            cid = lib.lookup(book_bits_key(src, start, r.pos))
            if cid is None:
                raise ValueError("codebook #%d not in game library (dims=%d entries=%d)" % (len(info.codebook_ids), de[0], de[1]))
            w.put(10, cid)
        info.codebook_ids.append(cid)

    # time domain transforms (Vorbis only: count-1=0, one 16-bit zero)
    if direction == W2V:
        w.put(6, 0); w.put(16, 0)
    else:
        tc = r.get(6) + 1
        for _ in range(tc):
            if r.get(16) != 0:
                raise ValueError("nonzero time domain value")

    # floors
    nfloors = f(6) + 1
    info.floor_count = nfloors
    for _ in range(nfloors):
        const(16, 1, "floor type")
        parts = f(5)
        classes = [f(4) for _ in range(parts)]
        maxc = max(classes) if classes else -1
        dims = []
        for _ in range(maxc + 1):
            dims.append(f(3) + 1)
            sub = f(2)
            if sub:
                f(8)
            for _ in range(1 << sub):
                f(8)
        f(2)
        rb = f(4)
        for c in classes:
            for _ in range(dims[c]):
                f(rb)

    # residues
    nres = f(6) + 1
    info.residue_count = nres
    for _ in range(nres):
        rtype = f(2, 16)
        if rtype > 2:
            raise ValueError("bad residue type")
        f(24); f(24); f(24)
        ncls = f(6) + 1
        f(8)
        cascade = []
        for _ in range(ncls):
            lo = f(3)
            hi = f(5) if f(1) else 0
            cascade.append(hi * 8 + lo)
        for c in cascade:
            for k in range(8):
                if c & (1 << k):
                    f(8)

    # mappings
    nmap = f(6) + 1
    info.mapping_count = nmap
    for _ in range(nmap):
        const(16, 0, "mapping type")
        submaps = f(4) + 1 if f(1) else 1
        if f(1):
            steps = f(8) + 1
            b = ilog(channels - 1)
            for _ in range(steps):
                f(b); f(b)
        if f(2) != 0:
            raise ValueError("mapping reserved != 0")
        if submaps > 1:
            for _ in range(channels):
                f(4)
        for _ in range(submaps):
            f(8); f(8); f(8)

    # modes
    nmodes = f(6) + 1
    for _ in range(nmodes):
        bf = f(1)
        const(16, 0, "window type")
        const(16, 0, "transform type")
        f(8)
        info.mode_blockflags.append(bf)

    if direction == W2V:
        w.put(1, 1)  # framing
    else:
        if r.get(1) != 1:
            raise ValueError("framing bit missing")
    return w.bytes(), info


def build_id_header(channels: int, rate: int, bs0_exp: int, bs1_exp: int) -> bytes:
    import struct
    return (b"\x01vorbis" + struct.pack("<IBIiii", 0, channels, rate, 0, 0, 0)
            + bytes([(bs1_exp << 4) | bs0_exp, 1]))


def build_comment_header(vendor: bytes = b"ackit") -> bytes:
    import struct
    return b"\x03vorbis" + struct.pack("<I", len(vendor)) + vendor + struct.pack("<I", 0) + b"\x01"


def ww_packet_to_vorbis(pkt: bytes, next_first: int, info: SetupInfo, prev_bf: int,
                         modified: bool = True) -> (bytes, int):
    """Rebuild a standard audio packet from a Wwise modified packet. Returns (packet, blockflag).
    When blocksizes are equal Wwise keeps standard packets (modified=False)."""
    mb = info.mode_bits
    if not modified:
        mode = (pkt[0] >> 1) & ((1 << mb) - 1)
        return pkt, info.mode_blockflags[mode]
    v = int.from_bytes(pkt, "little")
    nbits = len(pkt) * 8
    mode = v & ((1 << mb) - 1)
    bf = info.mode_blockflags[mode]
    out = 0  # packet type bit 0
    pos = 1
    out |= mode << pos; pos += mb
    if bf:
        nbf = 0 if next_first is None else info.mode_blockflags[next_first & ((1 << mb) - 1)]
        out |= prev_bf << pos; pos += 1
        out |= nbf << pos; pos += 1
    rest = v >> mb
    out |= rest << pos
    pos += nbits - mb
    return out.to_bytes((pos + 7) // 8, "little"), bf


def vorbis_packet_to_ww(pkt: bytes, info: SetupInfo, modified: bool = True) -> (bytes, int):
    """Strip packet-type bit and window flags. Returns (packet, blockflag)."""
    mb = info.mode_bits
    if not modified:
        if pkt[0] & 1:
            raise ValueError("not an audio packet")
        return pkt, info.mode_blockflags[(pkt[0] >> 1) & ((1 << mb) - 1)]
    v = int.from_bytes(pkt, "little")
    nbits = len(pkt) * 8
    if v & 1:
        raise ValueError("not an audio packet")
    mode = (v >> 1) & ((1 << mb) - 1)
    bf = info.mode_blockflags[mode]
    drop = 1 + mb + (2 if bf else 0)
    rest = v >> drop
    out = mode | (rest << mb)
    obits = nbits - (drop - mb)
    return out.to_bytes((obits + 7) // 8, "little"), bf