"""Encode audio into game-compatible WEMs (portable: libvorbis via libsndfile, no Wwise, no ffmpeg).

Pipeline: float32 PCM -> libvorbis (Ogg) -> `builder.ogg_to_wem` (Wwise 2011.1 layout, game codebook
IDs, FNV-1 hash, decode-alloc size). The quality is chosen so that every codebook libvorbis uses exists
in the game's built-in library (docs/formats/wem.md); setups the game itself ships are preferred.
"""
import io
import json
import os

import numpy as np
import soundfile as sf

from . import audio
from .builder import ogg_to_wem
from .convert import library
from .known_setups import lookup_setup
from .ogg import read_packets
from .paths import user_dir
from .wwvorbis import V2W, transcode_setup

QUALITIES = [q / 2 for q in range(20, -1, -1)]  # 10.0 .. 0.0 (libvorbis q; libsndfile cannot do q<0)
AUTO_MAX_Q_KNOWN = 6.0   # highest q considered when looking for a setup the game ships
AUTO_MAX_Q = 4.0         # otherwise, highest compatible q (game files sit around q3-5)
FALLBACK_RATES = (32000, 28000, 24000)          # e.g. 6 channels at 44.1/48 kHz
FIT_RATES = (24000, 22050, 16000, 12000, 11025, 8000)


class EncodeError(ValueError):
    pass


def vorbis_ogg(x, rate, quality):
    """Encode float32 PCM to Ogg Vorbis bytes at libvorbis quality `quality` (0..10)."""
    b = io.BytesIO()
    level = min(1.0, max(0.0, 1.0 - quality / 10.0))
    x = np.ascontiguousarray(audio._as2d(x))
    # write in blocks: libsndfile's Vorbis writer overflows the C stack on very large single writes
    with sf.SoundFile(b, "w", rate, x.shape[1], format="OGG", subtype="VORBIS", compression_level=level) as f:
        for i in range(0, len(x), 4096):
            f.write(x[i:i + 4096])
    return b.getvalue()


# --- setup probing (which q values are usable at a rate/channel count) ------------------------
_cache = None


def _cache_path():
    return os.path.join(user_dir("cache"), "vorbis_probe.json")


def _probe_cache():
    global _cache
    if _cache is None:
        try:
            c = json.load(open(_cache_path()))
            _cache = c if c.get("_lib") == sf.__libsndfile_version__ else None
        except (OSError, ValueError):
            _cache = None
        _cache = _cache or {"_lib": sf.__libsndfile_version__}
    return _cache


def probe_setup(rate, channels, quality):
    """(stripped Wwise setup or None, error) for libvorbis at these settings (cached)."""
    key = "%d,%d,%.1f" % (rate, channels, quality)
    c = _probe_cache()
    if key not in c:
        noise = (np.random.default_rng(7).standard_normal((rate, channels)) * 0.1).astype(np.float32)
        pk = read_packets(vorbis_ogg(noise, rate, quality))
        try:
            s, _ = transcode_setup(pk[2][0], V2W, channels, library())
            c[key] = [s.hex(), None]
        except ValueError as e:
            c[key] = [None, str(e)]
        try:
            json.dump(c, open(_cache_path(), "w"))
        except OSError:
            pass
    s, err = c[key]
    return (bytes.fromhex(s) if s else None), err


def choose_quality(rate, channels, quality="auto"):
    """-> (quality, known) for a rate/channel count, or raise EncodeError."""
    if quality != "auto":
        q = float(quality)
        s, err = probe_setup(rate, channels, q)
        if s is None:
            raise EncodeError("quality %s cannot be used at %d Hz / %d ch: %s" % (quality, rate, channels, err))
        return q, bool(lookup_setup(s))
    compatible = None
    for q in [x for x in QUALITIES if x <= AUTO_MAX_Q_KNOWN]:
        s, _ = probe_setup(rate, channels, q)
        if s is None:
            continue
        if lookup_setup(s):
            return q, True
        if compatible is None and q <= AUTO_MAX_Q:
            compatible = q
    if compatible is None:
        raise EncodeError("no libvorbis quality yields game-compatible codebooks at %d Hz / %d ch" % (rate, channels))
    return compatible, False


def pick_format(rate, channels, quality="auto"):
    """Find a usable (rate, quality). Falls back to lower rates when needed (Wwise resamples every
    source, so playback speed is unaffected)."""
    try:
        return (rate,) + choose_quality(rate, channels, quality)
    except EncodeError as first:
        for r in FALLBACK_RATES:
            if r < rate:
                try:
                    return (r,) + choose_quality(r, channels, quality)
                except EncodeError:
                    pass
        raise first


def encode(x, src_rate, rate, channels, quality="auto"):
    """PCM -> WEM bytes in the requested format. Returns (wem, info)."""
    r, q, known = pick_format(rate, channels, quality)
    pcm = audio.conform(x, src_rate, r, channels)
    if len(pcm) < 64:
        raise EncodeError("the audio is empty")
    wem = ogg_to_wem(vorbis_ogg(pcm, r, q))
    return wem, dict(rate=r, quality=q, known_setup=known, size=len(wem), channels=channels,
                     seconds=len(pcm) / float(r), requested_rate=rate)


def encode_fit(x, src_rate, rate, channels, limit=None, quality="auto", progress=None):
    """Like `encode`, but if the result is larger than `limit` bytes lower the quality step by step,
    then the sample rate, until it fits. Raises EncodeError if it cannot fit."""
    wem, info = encode(x, src_rate, rate, channels, quality)
    if limit is None or len(wem) <= limit:
        return wem, info
    first = len(wem)
    for r in [info["rate"]] + [f for f in FIT_RATES if f < info["rate"]]:
        pcm = audio.conform(x, src_rate, r, channels)
        for q in QUALITIES:
            if r == info["rate"] and q >= info["quality"]:
                continue
            if probe_setup(r, channels, q)[0] is None:
                continue
            if progress:
                progress("shrinking: %d Hz, quality %.1f" % (r, q))
            w = ogg_to_wem(vorbis_ogg(pcm, r, q))
            if len(w) <= limit:
                info.update(rate=r, quality=q, size=len(w), shrunk_from=first,
                            known_setup=bool(lookup_setup(probe_setup(r, channels, q)[0])))
                return w, info
    raise EncodeError("the audio is too long for this sound even at the lowest quality "
                      "(%d bytes, maximum %d). Shorten it." % (first, limit))