"""Build an Arkham City compatible WEM from a standard Ogg Vorbis stream.

Field rules were derived from the WEMs of the ITA WAD and are validated by
rebuilding original files byte-for-byte (`ackit verify`, tests/test_realgame.py).
"""
import struct

from .ogg import read_packets
from .wem import Wem, Chunk
from .wwvorbis import V2W, transcode_setup, vorbis_packet_to_ww
from .convert import packet_samples, library
from .known_setups import lookup_setup
from .allocsize import fnv1_32, decode_alloc

SEEK_INTERVAL = 8192


def build_seek_table(pkt_offsets, pkt_end_granules, interval=SEEK_INTERVAL):
    """pkt_offsets: offset of each audio packet relative to the setup start (= post seek-table)."""
    out = []
    prev_g, prev_o = 0, 0
    target = interval
    for off, g in zip(pkt_offsets, pkt_end_granules):
        if g >= target:
            out.append((g - prev_g, off - prev_o))
            prev_g, prev_o = g, off
            target = g + interval
    return out


def build_wem(channels, rate, bs0_exp, bs1_exp, ww_setup, ww_packets, blockflags, total_frames,
              codebook_hash, decode_alloc, decode_alloc_x64, channel_mask, seek=True, seek_interval=SEEK_INTERVAL, valid_bits=0x18,
              extra_chunks=()):
    samples = packet_samples(blockflags, 1 << bs0_exp, 1 << bs1_exp)
    cum = []
    g = 0
    for s in samples:
        g += s
        cum.append(g)
    last_extra = g - total_frames
    if not 0 <= last_extra < (1 << bs1_exp):
        raise ValueError("inconsistent total_frames (extra=%d)" % last_extra)

    setup_blob = struct.pack("<H", len(ww_setup)) + ww_setup
    offs = []
    pos = len(setup_blob)
    audio = bytearray()
    for p in ww_packets:
        if len(p) > 0xFFFF:
            raise ValueError("packet too large")
        offs.append(pos)
        audio += struct.pack("<H", len(p)) + p
        pos += 2 + len(p)
    seek_entries = build_seek_table(offs, cum, seek_interval) if seek else []
    for dg, do in seek_entries:
        if dg > 0xFFFF or do > 0xFFFF:
            raise ValueError("seek delta overflow")
    seek_blob = b"".join(struct.pack("<HH", a, b) for a, b in seek_entries)
    data = seek_blob + setup_blob + bytes(audio)
    st = len(seek_blob)

    w = Wem([Chunk(b"fmt ", b"")] + list(extra_chunks) + [Chunk(b"data", data)])
    w.channels = channels
    w.rate = rate
    w.avg_bps = (len(data) - st) * rate // total_frames
    w.valid_bits = valid_bits
    w.channel_mask = channel_mask
    w.vorb = dict(total_frames=total_frames,
                  loop_start_pkt=len(setup_blob),
                  loop_end_pkt=len(data) - st,
                  loop_begin_extra=0,
                  loop_end_extra=last_extra,
                  seek_table_size=st,
                  vorbis_data_offset=st + len(setup_blob),
                  max_packet_size=max(len(p) for p in ww_packets),
                  last_granule_extra=last_extra,
                  decode_alloc=decode_alloc,
                  decode_alloc_x64=decode_alloc_x64,
                  codebook_hash=codebook_hash,
                  bs0_exp=bs0_exp, bs1_exp=bs1_exp)
    return w.to_bytes()


DEFAULT_MASKS = {1: 0x4, 2: 0x3, 4: 0x33, 5: 0x37, 6: 0x3F}


def ogg_to_wem(ogg: bytes, total_frames=None, allow_unknown_setup=True):
    """Convert a libvorbis Ogg stream into a game-compatible WEM (bytes)."""
    lib = library()
    pk = read_packets(ogg)
    idh = pk[0][0]
    if idh[:7] != b"\x01vorbis":
        raise ValueError("not Ogg Vorbis")
    channels = idh[11]
    rate = struct.unpack_from("<I", idh, 12)[0]
    bs0_exp, bs1_exp = idh[28] & 15, idh[28] >> 4
    ww_setup, info = transcode_setup(pk[2][0], V2W, channels, lib)
    ww_packets, flags = [], []
    last_gran = None
    for p, g in pk[3:]:
        if not p:
            continue
        wp, bf = vorbis_packet_to_ww(p, info, bs0_exp != bs1_exp)
        ww_packets.append(wp)
        flags.append(bf)
        if g is not None:
            last_gran = g
    if total_frames is None:
        total_frames = last_gran
    hdr = setup_header_values(ww_setup, channels, allow_unknown_setup)
    return build_wem(channels, rate, bs0_exp, bs1_exp, ww_setup, ww_packets, flags, total_frames,
                     hdr["hash"], hdr["alloc"], hdr["alloc64"],
                     DEFAULT_MASKS.get(channels, (1 << channels) - 1))


def setup_header_values(ww_setup: bytes, channels: int, allow_unknown_setup=True) -> dict:
    """hash/alloc values for a setup: from the game's own files when known, else computed
    with the reversed formulas (identical results on all known setups)."""
    known = lookup_setup(ww_setup)
    h = fnv1_32(ww_setup)
    alloc, alloc64 = decode_alloc(ww_setup, channels, library())
    if known is not None:
        if (known["hash"], known["alloc"]) != (h, alloc):
            raise AssertionError("formula mismatch on known setup")
        return dict(hash=known["hash"], alloc=known["alloc"], alloc64=known["alloc64"], known=True)
    if not allow_unknown_setup:
        raise ValueError("setup not used by any game file")
    return dict(hash=h, alloc=alloc, alloc64=alloc64, known=False)