# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
#
# The log-mel feature extraction re-implements the defaults of Gemma4AudioFeatureExtractor
# (feature_extraction_gemma4.py in Hugging Face transformers 5.19, Copyright 2026 Google LLC, Apache-2.0) as used by
# sentence-transformers for google/embeddinggemma-2, and the HTK mel filter-bank formula of transformers'
# audio_utils.py (Copyright 2023 The HuggingFace Inc. team and the librosa & torchaudio authors, Apache-2.0; the
# formula as in librosa / torchaudio). No transformers code is included.
"""Audio front end for EmbeddingGemma 2 on Core AI: file -> 16 kHz mono -> log-mel frames -> model inputs (numpy only).

Reference pipeline (what upstream does with a file path: librosa.load(sr=16000) then the feature extractor):
  1. decode to float32 at the file's own rate; 2. mix to mono (mean of the channels, float32);
  3. resample to 16 kHz; 4. keep the first 30 s (480 000 samples); 5. log-mel frames every 10 ms.
This module keeps 1, 2, 4 and 5 identical to upstream. For 3, upstream uses libsoxr "HQ"; this package uses its own
Kaiser-windowed sinc (polyphase, 125 dB stop band, pass band 0.913 of 8 kHz), implemented identically here and in
Swift (`AudioResampler`), see README "Audio: decoding and resampling" for the measured difference (16 kHz input,
the model's native rate, is never resampled and is exact).

    import eg2_audio_frontend
    wav = eg2_audio_frontend.load_mono_16k("clip.wav")               # float32 [n] at 16 kHz
    inputs = eg2_audio_frontend.prepare_audio_inputs(wav)           # AudioInputs for the audio_<Tb> function
    ids = eg2_audio_frontend.audio_sequence(inputs.n_tokens, ...)   # <bos> <boa> [n] <eoa> <eos>

Decoders: `soundfile` (libsndfile, the decoder upstream's librosa uses: WAV, AIFF, FLAC, OGG, CAF/ALAC, MP3 with
gapless trimming) when it is installed, otherwise, and for AAC/M4A always, macOS `afconvert` (Apple's AudioToolbox
decoder; MP3 files are trimmed with the LAME/Info tag here so that they match libsndfile; see README).
"""
from __future__ import annotations

import math
import os
import shutil
import struct
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path

import numpy as np

SAMPLE_RATE = 16000
FRAME_LENGTH = 320          # 20 ms
HOP_LENGTH = 160            # 10 ms
FFT_LENGTH = 512
N_MELS = 128
F_MIN, F_MAX = 0.0, 8000.0
MEL_FLOOR = 1e-3
MAX_SAMPLES = 480_000       # 30 s: the upstream default (truncation=True)
PAD_MULTIPLE = 128          # samples
FRAME_BUCKETS = [192, 384, 1024, 1536, 3000]
MAX_FRAMES = 2999           # valid frames of 30 s

# resampler design (shared with Swift AudioResampler; do not change one without the other)
RESAMPLE_PASSBAND = 0.913   # of the lower Nyquist (8 kHz when downsampling)
RESAMPLE_STOPBAND = 1.0
RESAMPLE_ATTEN_DB = 125.0


# ------------------------------------------------------------------------------------------------ decoding
@dataclass
class DecodedAudio:
    samples: np.ndarray      # float32 [frames, channels]
    sample_rate: int
    decoder: str


SOUNDFILE_EXT = {".wav", ".wave", ".aif", ".aiff", ".aifc", ".flac", ".ogg", ".oga", ".mp3", ".caf", ".w64", ".rf64"}


def decode(path: str | os.PathLike) -> DecodedAudio:
    """float32 [frames, channels] at the file's own rate. libsndfile (soundfile) where available, Apple's decoder
    (afconvert) otherwise; MP3 encoder delay/padding trimmed in both cases (libsndfile does it itself)."""
    path = Path(path)
    if path.suffix.lower() in SOUNDFILE_EXT:
        try:
            import soundfile as sf
        except ImportError:
            sf = None
        if sf is not None:
            try:
                x, sr = sf.read(str(path), dtype="float32", always_2d=True)
                return DecodedAudio(np.ascontiguousarray(x, dtype=np.float32), int(sr), "libsndfile")
            except Exception:  # noqa: BLE001 - fall back to Apple's decoder
                pass
    return _decode_afconvert(path)


def _decode_afconvert(path: Path) -> DecodedAudio:
    afconvert = shutil.which("afconvert")
    if not afconvert:
        raise RuntimeError(f"cannot decode {path.name}: install soundfile (pip install soundfile) or run on macOS")
    with tempfile.TemporaryDirectory() as td:
        out = Path(td) / "decoded.wav"
        subprocess.run([afconvert, "-f", "WAVE", "-d", "LEF32", str(path), str(out)], check=True, capture_output=True)
        x, sr = _read_wav_f32(out)
    if path.suffix.lower() == ".mp3":
        x = _mp3_trim(x, mp3_gapless_info(path))
    return DecodedAudio(x, sr, "afconvert")


def _mp3_trim(x: np.ndarray, info: tuple[int, int, int, int] | None) -> np.ndarray:
    """Brings Apple's MP3 output to libsndfile's gapless length. AudioToolbox applies the LAME tag itself for most
    files, returns whole frames for some, and for others trims the delay but keeps one or two whole frames too many
    at the end (it counts the Xing/Info frame as audio: `afinfo` shows one packet more than the tag's frame count).
    Same rule in Swift (AudioDecoder)."""
    if info is None:
        return x
    frames, spf, delay, padding = info
    full = frames * spf
    gapless = full - delay - padding
    if len(x) == full:                                   # untrimmed: cut delay and padding
        return x[delay: len(x) - padding] if padding else x[delay:]
    if gapless > 0 and len(x) > gapless and (len(x) - gapless) % spf == 0 and len(x) - gapless <= 2 * spf:
        return x[:gapless]                               # delay trimmed, whole extra frames at the end
    return x


def _read_wav_f32(path: Path) -> tuple[np.ndarray, int]:
    b = path.read_bytes()
    assert b[:4] == b"RIFF" and b[8:12] == b"WAVE", "afconvert did not write a WAVE file"
    p, fmt, data = 12, None, None
    while p + 8 <= len(b):
        cid, size = b[p:p + 4], struct.unpack("<I", b[p + 4:p + 8])[0]
        body = b[p + 8:p + 8 + size]
        if cid == b"fmt ":
            fmt = struct.unpack("<HHIIHH", body[:16])
        elif cid == b"data":
            data = body
        p += 8 + size + (size & 1)
    assert fmt is not None and data is not None
    tag, ch, sr, _, _, bits = fmt
    assert bits == 32 and tag in (3, 0xFFFE), "expected float32 WAVE"
    return np.frombuffer(data, dtype="<f4").reshape(-1, ch).copy(), sr


def mp3_gapless_info(path: Path) -> tuple[int, int, int, int] | None:
    """(frames, samples per frame, encoder delay, padding) from the Xing/Info + LAME tag of the first MPEG frame,
    None if absent. Apple's decoder compensates the 529-sample decoder delay itself; libsndfile (mpg123) trims
    the delay and the padding too (gapless), which is the reference behaviour."""
    with open(path, "rb") as f:
        b = f.read(1 << 20)
    i = 0
    if b[:3] == b"ID3":
        i = 10 + (((b[6] & 0x7F) << 21) | ((b[7] & 0x7F) << 14) | ((b[8] & 0x7F) << 7) | (b[9] & 0x7F))
    while i < len(b) - 4 and not (b[i] == 0xFF and (b[i + 1] & 0xE0) == 0xE0):
        i += 1
    if i >= len(b) - 4:
        return None
    mpeg1 = ((b[i + 1] >> 3) & 3) == 3
    mono = ((b[i + 3] >> 6) & 3) == 3
    side = (17 if mono else 32) if mpeg1 else (9 if mono else 17)
    p = i + 4 + side
    if b[p:p + 4] not in (b"Xing", b"Info"):
        return None
    flags = struct.unpack(">I", b[p + 4:p + 8])[0]
    p += 8
    if not flags & 1:
        return None
    frames = struct.unpack(">I", b[p:p + 4])[0]
    p += 4 + (4 if flags & 2 else 0) + (100 if flags & 4 else 0) + (4 if flags & 8 else 0)
    d = b[p + 21:p + 24]
    if len(d) < 3 or not b[p:p + 4].isalpha():
        return None
    return frames, 1152 if mpeg1 else 576, (d[0] << 4) | (d[1] >> 4), ((d[1] & 0x0F) << 8) | d[2]


def to_mono(x: np.ndarray) -> np.ndarray:
    """Mean over channels in float32 (librosa.to_mono), [frames, channels] -> [frames]."""
    x = np.asarray(x, dtype=np.float32)
    if x.ndim == 1:
        return x
    if x.shape[1] == 1:
        return np.ascontiguousarray(x[:, 0])
    acc = x[:, 0].copy()
    for c in range(1, x.shape[1]):   # sequential float32 sum, then one divide (same order as Swift)
        acc = acc + x[:, c]
    return (acc / np.float32(x.shape[1])).astype(np.float32)


# ------------------------------------------------------------------------------------------------ resampling
def _bessel_i0(x: float) -> float:
    """Modified Bessel function of order 0 by its power series (same code in Swift)."""
    s, t, k = 1.0, 1.0, 1
    h = x / 2.0
    while True:
        t *= (h / k) ** 2
        s += t
        if t < s * 1e-17:
            return s
        k += 1


def sinc_filter(up: int, down: int, in_rate: int) -> tuple[np.ndarray, int]:
    """Kaiser-windowed sinc low-pass for the polyphase resampler in_rate -> in_rate*up/down (float64 taps, odd length,
    unity DC gain). Returns (taps, delay = (len-1)/2). Pass band 0.913 and stop band 1.0 of the lower Nyquist."""
    nyq_low = min(in_rate, in_rate * up // down) / 2.0
    nyq_hi = in_rate * up / 2.0
    fc = 0.5 * (RESAMPLE_PASSBAND + RESAMPLE_STOPBAND) * nyq_low / nyq_hi    # normalized cutoff (1 = nyq_hi)
    tw = (RESAMPLE_STOPBAND - RESAMPLE_PASSBAND) * nyq_low / nyq_hi
    beta = 0.1102 * (RESAMPLE_ATTEN_DB - 8.7)
    n = int(math.ceil((RESAMPLE_ATTEN_DB - 7.95) / (2.285 * math.pi * tw)))
    n |= 1
    c = (n - 1) // 2
    i0b = _bessel_i0(beta)
    h = np.empty(n, dtype=np.float64)
    for k in range(n):
        m = k - c
        sinc = fc if m == 0 else math.sin(math.pi * fc * m) / (math.pi * m)
        r = 2.0 * k / (n - 1) - 1.0
        w = _bessel_i0(beta * math.sqrt(max(0.0, 1.0 - r * r))) / i0b
        h[k] = sinc * w
    h /= h.sum()
    return h, c


_FILTERS: dict[tuple[int, int, int], tuple[np.ndarray, int]] = {}


def resample_16k(x: np.ndarray, sr: int) -> np.ndarray:
    """float32 [n] at sr -> float32 [ceil(n*16000/sr)] at 16 kHz (output length as librosa's fix=True).
    Polyphase windowed sinc computed in float64: y[j] = up * sum_m x[m] h[c + j*down - m*up]."""
    x = np.asarray(x, dtype=np.float32)
    if sr == SAMPLE_RATE:
        return x
    n_out = int(math.ceil(len(x) * SAMPLE_RATE / sr))
    g = math.gcd(SAMPLE_RATE, sr)
    up, down = SAMPLE_RATE // g, sr // g
    key = (up, down, sr)
    if key not in _FILTERS:
        _FILTERS[key] = sinc_filter(up, down, sr)
    h, c = _FILTERS[key]
    xd = x.astype(np.float64)
    y = np.zeros(n_out, dtype=np.float64)
    nh = len(h)
    # polyphase branches: outputs j with the same phase p = (j*down) % up share the taps h[c + p + r*up]
    for p in range(up):
        # outputs j with (j*down) % up == p
        j0 = next((j for j in range(up) if (j * down) % up == p), None)
        if j0 is None:
            continue
        js = np.arange(j0, n_out, up)
        if len(js) == 0:
            continue
        # for these outputs, c + j*down - m*up = c + p + q*up - m*up with q = (j*down - p)//up  -> taps at offsets
        # k = c + p + (q - m)*up ; let r = q - m, k in [0, nh): r from ceil((-c-p)/up) to floor((nh-1-c-p)/up)
        r_min = -((c + p) // up)
        r_max = (nh - 1 - c - p) // up
        q = (js * down - p) // up
        acc = np.zeros(len(js), dtype=np.float64)
        for r in range(r_min, r_max + 1):
            k = c + p + r * up
            m = q - r
            ok = (m >= 0) & (m < len(xd))
            if not ok.any():
                continue
            acc[ok] += h[k] * xd[m[ok]]
        y[js] = acc * up
    return y.astype(np.float32)


def load_mono_16k(path: str | os.PathLike) -> np.ndarray:
    """decode -> mono -> 16 kHz (the librosa.load(sr=16000, mono=True) order), float32 [n]."""
    d = decode(path)
    return resample_16k(to_mono(d.samples), d.sample_rate)


# ------------------------------------------------------------------------------------------------ features
def _hz_to_mel_htk(f):
    return 2595.0 * np.log10(1.0 + np.asarray(f, dtype=np.float64) / 700.0)


def _mel_to_hz_htk(m):
    return 700.0 * (10.0 ** (np.asarray(m, dtype=np.float64) / 2595.0) - 1.0)


def mel_filters() -> np.ndarray:
    """[257, 128] float64 triangular HTK filters, unnormalized (transformers mel_filter_bank(norm=None, mel_scale='htk'))."""
    n_bins = FFT_LENGTH // 2 + 1
    mels = np.linspace(_hz_to_mel_htk(F_MIN), _hz_to_mel_htk(F_MAX), N_MELS + 2)
    filter_freqs = _mel_to_hz_htk(mels)
    fft_freqs = np.linspace(0, SAMPLE_RATE // 2, n_bins)
    diff = np.diff(filter_freqs)
    slopes = filter_freqs[None, :] - fft_freqs[:, None]
    down = -slopes[:, :-2] / diff[:-1]
    up = slopes[:, 2:] / diff[1:]
    return np.maximum(0.0, np.minimum(down, up))


def hann_periodic(n: int = FRAME_LENGTH) -> np.ndarray:
    return (0.5 - 0.5 * np.cos(2 * np.pi * np.arange(n) / n)).astype(np.float32)


_MEL = None
_WIN = None


def num_valid_frames(n_samples: int) -> int:
    """Frames whose whole 321-sample analysis window (incl. the 160-sample left pad) lies in real audio (max 2999)."""
    n = min(n_samples, MAX_SAMPLES)
    return max(0, (n + FRAME_LENGTH // 2 - (FRAME_LENGTH + 1)) // HOP_LENGTH + 1)


def n_audio_tokens(n_valid_frames: int) -> int:
    return ((n_valid_frames + 1) // 2 + 1) // 2


def log_mel(wav16k: np.ndarray, max_samples: int | None = MAX_SAMPLES) -> np.ndarray:
    """[n_valid_frames, 128] float32 log-mel features, bit-for-bit the upstream feature extractor (invalid frames,
    which upstream zeroes and masks, are not returned)."""
    global _MEL, _WIN
    if _MEL is None:
        _MEL, _WIN = mel_filters(), hann_periodic()
    x = np.asarray(wav16k, dtype=np.float32)
    if max_samples is not None:
        x = x[:max_samples]
    n = len(x)
    T = num_valid_frames(n)
    if T <= 0:
        return np.zeros((0, N_MELS), np.float32)
    n_pad = int(math.ceil(n / PAD_MULTIPLE) * PAD_MULTIPLE)
    xp = np.zeros(FRAME_LENGTH // 2 + n_pad, dtype=np.float32)
    xp[FRAME_LENGTH // 2: FRAME_LENGTH // 2 + n] = x
    idx = np.arange(T)[:, None] * HOP_LENGTH + np.arange(FRAME_LENGTH)[None, :]
    frames = xp[idx] * _WIN                                   # float32
    mag = np.abs(np.fft.rfft(frames, n=FFT_LENGTH, axis=-1))  # complex64 -> float32 (numpy >= 2)
    mel = mag @ _MEL                                          # float64
    return np.log(mel + MEL_FLOOR).astype(np.float32)


# ------------------------------------------------------------------------------------------------ model inputs
@dataclass
class AudioInputs:
    features: np.ndarray     # [1, Tb, 128] float16, zero beyond n_frames
    frame_mask: np.ndarray   # [1, Tb] float16
    half_mask: np.ndarray    # [1, Tb/2] float16
    token_mask: np.ndarray   # [1, Tb/4] float16
    n_frames: int
    n_tokens: int
    bucket: int


def bucket_for(n_frames: int, buckets=FRAME_BUCKETS) -> int:
    for b in buckets:
        if b >= n_frames:
            return b
    raise ValueError(f"{n_frames} frames exceed the largest bucket {buckets[-1]}")


def prepare_audio_inputs(wav16k: np.ndarray, buckets=FRAME_BUCKETS, max_samples: int | None = MAX_SAMPLES) -> AudioInputs:
    """16 kHz mono float32 -> inputs of audio_<bucket>. Shorter than 161 samples (0 valid frames) gives 0 tokens
    (upstream crashes below 129 samples; here any length is accepted)."""
    f = log_mel(wav16k, max_samples)
    nv = f.shape[0]
    Tb = bucket_for(max(nv, 1), buckets)
    x = np.zeros((1, Tb, N_MELS), np.float16)
    x[0, :nv] = f
    fm = np.zeros((1, Tb), np.float16); fm[0, :nv] = 1
    hm = np.zeros((1, Tb // 2), np.float16); hm[0, :(nv + 1) // 2] = 1
    nt = n_audio_tokens(nv)
    tm = np.zeros((1, Tb // 4), np.float16); tm[0, :nt] = 1
    return AudioInputs(x, fm, hm, tm, nv, nt, Tb)


def audio_sequence(n_tokens: int, bos: int, boa: int, audio: int, eoa: int, eos: int) -> list[int]:
    """The upstream single-clip sequence (no prompt): <bos> <|audio> <|audio|> x n <audio|> <eos>."""
    return [bos, boa] + [audio] * n_tokens + [eoa, eos]


def chunks_30s(wav16k: np.ndarray, chunk_samples: int = MAX_SAMPLES) -> list[np.ndarray]:
    """Consecutive 30 s windows (the last one shorter) for the optional chunk-mean embedding of long audio. A trailing
    window too short for one frame (< 161 samples, 0 tokens) is dropped unless it is the only one."""
    x = np.asarray(wav16k, dtype=np.float32)
    parts = [x[i:i + chunk_samples] for i in range(0, max(len(x), 1), chunk_samples)]
    return [p for i, p in enumerate(parts) if i == 0 or len(p) > FRAME_LENGTH // 2]