# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

import base64
import binascii
import io
import math
import wave

import numpy as np

from vllm_omni.utils.audio_resample import StreamingAudioResampler

try:
    from audioop import alaw2lin, lin2alaw, lin2ulaw, ulaw2lin
except ImportError:  # pragma: no cover - audioop is removed in newer Python.
    alaw2lin = lin2alaw = lin2ulaw = ulaw2lin = None


MIN_INPUT_SAMPLE_RATE_HZ = 8_000
MAX_INPUT_SAMPLE_RATE_HZ = 192_000


def validate_input_sample_rate_hz(sample_rate_hz: object) -> int:
    if (
        isinstance(sample_rate_hz, bool)
        or not isinstance(sample_rate_hz, int | float)
        or (isinstance(sample_rate_hz, float) and not math.isfinite(sample_rate_hz))
    ):
        raise ValueError("sample_rate_hz must be a finite integer")
    rate = int(sample_rate_hz)
    if rate != sample_rate_hz:
        raise ValueError("sample_rate_hz must be an integer")
    if not MIN_INPUT_SAMPLE_RATE_HZ <= rate <= MAX_INPUT_SAMPLE_RATE_HZ:
        raise ValueError(f"sample_rate_hz must be between {MIN_INPUT_SAMPLE_RATE_HZ} and {MAX_INPUT_SAMPLE_RATE_HZ}")
    return rate


def resample_pcm16_mono(raw: bytes, *, source_rate_hz: int, target_rate_hz: int) -> bytes:
    if source_rate_hz <= 0 or target_rate_hz <= 0 or source_rate_hz == target_rate_hz:
        return raw
    samples = np.frombuffer(raw, dtype="<i2")
    rate_gcd = math.gcd(source_rate_hz, target_rate_hz)
    reduced_factor = max(source_rate_hz // rate_gcd, target_rate_hz // rate_gcd)
    if reduced_factor > StreamingAudioResampler.max_polyphase_factor:
        # The FIR size scales with the reduced ratio. Keep unusual client
        # rates bounded while common audio rates use the high-quality path.
        if samples.size <= 1:
            return raw
        target_size = max(1, int(round(samples.size * target_rate_hz / source_rate_hz)))
        source_x = np.linspace(0.0, 1.0, num=samples.size, endpoint=True)
        target_x = np.linspace(0.0, 1.0, num=target_size, endpoint=True)
        resampled = np.interp(target_x, source_x, samples)
        return np.clip(resampled, -32_768, 32_767).astype("<i2").tobytes()
    resampler = StreamingAudioResampler(source_rate_hz, target_rate_hz)
    normalized = samples.astype(np.float32) * np.float32(1.0 / 32768.0)
    resampled = resampler.process(normalized, final=True)
    return np.clip(np.rint(resampled * 32768.0), -32768, 32767).astype("<i2").tobytes()


def decode_g711_ulaw(raw: bytes) -> bytes:
    if ulaw2lin is not None:
        return ulaw2lin(raw, 2)
    data = np.frombuffer(raw, dtype=np.uint8)
    value = np.bitwise_not(data).astype(np.int16)
    sign = value & 0x80
    exponent = (value >> 4) & 0x07
    mantissa = value & 0x0F
    sample = ((mantissa << 3) + 0x84) << exponent
    return np.where(sign != 0, -(sample - 0x84), sample - 0x84).astype("<i2").tobytes()


def decode_g711_alaw(raw: bytes) -> bytes:
    if alaw2lin is not None:
        return alaw2lin(raw, 2)
    data = np.bitwise_xor(np.frombuffer(raw, dtype=np.uint8), 0x55).astype(np.int16)
    sign = data & 0x80
    exponent = (data >> 4) & 0x07
    mantissa = data & 0x0F
    sample = np.where(exponent == 0, (mantissa << 4) + 8, ((mantissa << 4) + 0x108) << (exponent - 1))
    return np.where(sign != 0, sample, -sample).astype("<i2").tobytes()


def encode_g711_ulaw(raw: bytes) -> bytes:
    if lin2ulaw is not None:
        return lin2ulaw(raw, 2)
    pcm = np.clip(np.frombuffer(raw, dtype="<i2").astype(np.int32), -32635, 32635)
    sign = np.where(pcm < 0, 0x80, 0)
    magnitude = np.abs(pcm) + 0x84
    exponent = np.zeros_like(magnitude)
    for exp in range(7):
        exponent = np.where(magnitude > (0xFF << exp), exp + 1, exponent)
    mantissa = (magnitude >> (exponent + 3)) & 0x0F
    return (np.bitwise_not(sign | (exponent << 4) | mantissa) & 0xFF).astype(np.uint8).tobytes()


def encode_g711_alaw(raw: bytes) -> bytes:
    if lin2alaw is not None:
        return lin2alaw(raw, 2)
    pcm = np.frombuffer(raw, dtype="<i2").astype(np.int32)
    sign = np.where(pcm >= 0, 0x80, 0x00)
    magnitude = np.abs(pcm)
    exponent = np.zeros_like(magnitude)
    for exp in range(1, 8):
        exponent = np.where(magnitude >= (1 << (exp + 7)), exp, exponent)
    mantissa = np.where(
        exponent == 0,
        (magnitude >> 4) & 0x0F,
        (magnitude >> (exponent + 3)) & 0x0F,
    )
    return ((sign | (exponent << 4) | mantissa) ^ 0x55).astype(np.uint8).tobytes()


def wav_payload_to_pcm16(raw: bytes) -> tuple[bytes | None, int | None]:
    try:
        with wave.open(io.BytesIO(raw), "rb") as wav_file:
            channels = wav_file.getnchannels()
            sample_width = wav_file.getsampwidth()
            sample_rate_hz = wav_file.getframerate()
            frames = wav_file.readframes(wav_file.getnframes())
        if sample_width != 2:
            return None, sample_rate_hz
        if channels <= 1:
            return frames, sample_rate_hz
        pcm = np.frombuffer(frames, dtype="<i2").reshape(-1, channels)
        mono = np.mean(pcm.astype(np.float32), axis=1)
        return np.clip(mono, -32768, 32767).astype("<i2").tobytes(), sample_rate_hz
    except (EOFError, ValueError, wave.Error):
        return None, None


def encode_float32_mono_wav_base64(
    samples: np.ndarray,
    *,
    sample_rate_hz: int,
) -> str:
    """Package normalized mono float32 samples as a PCM16 WAV data payload."""
    rate = validate_input_sample_rate_hz(sample_rate_hz)
    normalized = np.ascontiguousarray(samples, dtype=np.float32).reshape(-1)
    pcm16 = np.clip(np.rint(normalized * 32768.0), -32768, 32767).astype("<i2")
    with io.BytesIO() as buffer:
        with wave.open(buffer, "wb") as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(rate)
            wav_file.writeframes(pcm16.tobytes())
        return base64.b64encode(buffer.getvalue()).decode("ascii")


def convert_input_audio_with_rate(
    audio: object,
    fmt: object,
    *,
    sample_rate_hz: int | float | None = None,
    target_sample_rate_hz: int = 16_000,
) -> tuple[object, object, int | float | None]:
    if not isinstance(audio, str) or not isinstance(fmt, str):
        return audio, fmt, sample_rate_hz
    normalized = fmt.lower()
    if normalized not in {"pcm16", "pcm_s16le", "s16le", "g711_ulaw", "g711_alaw"}:
        return audio, fmt, sample_rate_hz
    if isinstance(sample_rate_hz, int | float):
        sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz)
    try:
        raw = base64.b64decode(audio.strip(), validate=False)
    except (binascii.Error, ValueError):
        return audio, fmt, sample_rate_hz
    if normalized == "g711_ulaw":
        raw = decode_g711_ulaw(raw)
        sample_rate_hz = sample_rate_hz if isinstance(sample_rate_hz, int | float) else 8_000
    elif normalized == "g711_alaw":
        raw = decode_g711_alaw(raw)
        sample_rate_hz = sample_rate_hz if isinstance(sample_rate_hz, int | float) else 8_000
    elif len(raw) % 2:
        return audio, fmt, sample_rate_hz
    if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) != target_sample_rate_hz:
        raw = resample_pcm16_mono(
            raw,
            source_rate_hz=int(sample_rate_hz),
            target_rate_hz=target_sample_rate_hz,
        )
        sample_rate_hz = target_sample_rate_hz
    pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
    encoded = base64.b64encode(np.ascontiguousarray(pcm, dtype="<f4").tobytes()).decode("ascii")
    return encoded, "pcm_f32le", sample_rate_hz


def convert_output_audio(
    audio: str,
    *,
    source_fmt: str,
    target_fmt: str,
    source_sample_rate_hz: int | None = None,
    target_sample_rate_hz: int | None = None,
) -> tuple[str, str, int | None]:
    target = target_fmt.lower()
    if target not in {"g711_ulaw", "g711_alaw"}:
        return audio, source_fmt, source_sample_rate_hz
    try:
        raw = base64.b64decode(audio, validate=False)
    except (binascii.Error, ValueError):
        return audio, source_fmt, source_sample_rate_hz
    source = source_fmt.lower()
    if source == "wav":
        pcm_raw, wav_rate = wav_payload_to_pcm16(raw)
        if pcm_raw is None:
            return audio, source_fmt, source_sample_rate_hz
        raw = pcm_raw
        source_sample_rate_hz = source_sample_rate_hz or wav_rate
    elif source not in {"pcm", "pcm16", "pcm_s16le", "s16le"} or len(raw) % 2:
        return audio, source_fmt, source_sample_rate_hz
    target_rate = target_sample_rate_hz or 8_000
    if source_sample_rate_hz is not None:
        raw = resample_pcm16_mono(raw, source_rate_hz=source_sample_rate_hz, target_rate_hz=target_rate)
    encoded = encode_g711_ulaw(raw) if target == "g711_ulaw" else encode_g711_alaw(raw)
    return base64.b64encode(encoded).decode("ascii"), target, target_rate