from __future__ import annotations

import audioop
import math
import random
import wave
from dataclasses import dataclass
from pathlib import Path


TARGET_SAMPLE_RATE = 16_000
TARGET_CHANNELS = 1
SAMPLE_WIDTH = 2
FRAME_MS = 20


@dataclass(frozen=True, slots=True)
class AudioSample:
    name: str
    pcm: bytes
    sample_rate: int = TARGET_SAMPLE_RATE
    channels: int = TARGET_CHANNELS


@dataclass(frozen=True, slots=True)
class AudioMetrics:
    duration_ms: int
    rms_dbfs: float
    peak_dbfs: float
    clipping_ratio: float
    bytes_len: int

    @property
    def audible(self) -> bool:
        return self.rms_dbfs > -55.0


@dataclass(frozen=True, slots=True)
class VADProxyMetrics:
    threshold_dbfs: float
    prefix_padding_ms: int
    min_speech_ms: int
    first_voice_ms: int
    last_voice_ms: int
    speech_ms: int
    padded_start_ms: int
    unrecovered_prefix_ms: int
    initial_200ms_dbfs: float
    initial_500ms_dbfs: float
    low_start_risk: bool


def _dbfs(value: float) -> float:
    if value <= 0:
        return -120.0
    return 20.0 * math.log10(value / 32768.0)


def _samples_from_pcm(pcm: bytes) -> list[int]:
    if not pcm:
        return []
    count = len(pcm) // SAMPLE_WIDTH
    return list(audioop.getsample(pcm, SAMPLE_WIDTH, idx) for idx in range(count))


def _pcm_from_samples(samples: list[int]) -> bytes:
    out = bytearray()
    for value in samples:
        clipped = max(-32768, min(32767, int(round(value))))
        out.extend(clipped.to_bytes(2, byteorder="little", signed=True))
    return bytes(out)


def _silence_pcm(sample_rate: int, duration_ms: int) -> bytes:
    return b"\x00" * round(sample_rate * duration_ms / 1000) * SAMPLE_WIDTH


def _split_pcm_at_ms(pcm: bytes, *, sample_rate: int, split_ms: int) -> tuple[bytes, bytes]:
    byte_count = max(0, round(sample_rate * split_ms / 1000) * SAMPLE_WIDTH)
    byte_count -= byte_count % SAMPLE_WIDTH
    return pcm[:byte_count], pcm[byte_count:]


def audio_metrics(pcm: bytes, *, sample_rate: int = TARGET_SAMPLE_RATE, channels: int = TARGET_CHANNELS) -> AudioMetrics:
    samples = _samples_from_pcm(pcm)
    if not samples:
        return AudioMetrics(duration_ms=0, rms_dbfs=-120.0, peak_dbfs=-120.0, clipping_ratio=0.0, bytes_len=0)

    rms = audioop.rms(pcm, SAMPLE_WIDTH)
    peak = max(abs(item) for item in samples)
    clipped = sum(1 for item in samples if abs(item) >= 32760)
    sample_count_per_channel = len(samples) / max(1, channels)
    duration_ms = round((sample_count_per_channel / max(1, sample_rate)) * 1000)
    return AudioMetrics(
        duration_ms=duration_ms,
        rms_dbfs=round(_dbfs(rms), 2),
        peak_dbfs=round(_dbfs(peak), 2),
        clipping_ratio=round(clipped / len(samples), 6),
        bytes_len=len(pcm),
    )


def _window_dbfs(
    pcm: bytes,
    *,
    sample_rate: int = TARGET_SAMPLE_RATE,
    duration_ms: int,
) -> float:
    byte_count = max(0, round(sample_rate * duration_ms / 1000) * SAMPLE_WIDTH)
    byte_count -= byte_count % SAMPLE_WIDTH
    window = pcm[:byte_count] if byte_count else b""
    if not window:
        return -120.0
    return round(_dbfs(audioop.rms(window, SAMPLE_WIDTH)), 2)


def vad_proxy_metrics(
    sample: AudioSample,
    *,
    threshold_dbfs: float = -45.0,
    prefix_padding_ms: int = 1000,
    min_speech_ms: int = 100,
    frame_ms: int = FRAME_MS,
) -> VADProxyMetrics:
    frame_bytes = max(1, round(sample.sample_rate * frame_ms / 1000)) * SAMPLE_WIDTH
    threshold_rms = int(32768 * (10 ** (threshold_dbfs / 20.0)))
    active_frames: list[int] = []
    total_frames = max(1, math.ceil(len(sample.pcm) / frame_bytes))
    for frame_idx in range(total_frames):
        frame = sample.pcm[frame_idx * frame_bytes : (frame_idx + 1) * frame_bytes]
        if len(frame) < frame_bytes:
            frame = frame + b"\x00" * (frame_bytes - len(frame))
        if audioop.rms(frame, SAMPLE_WIDTH) >= threshold_rms:
            active_frames.append(frame_idx)

    initial_200ms_dbfs = _window_dbfs(sample.pcm, sample_rate=sample.sample_rate, duration_ms=200)
    initial_500ms_dbfs = _window_dbfs(sample.pcm, sample_rate=sample.sample_rate, duration_ms=500)
    if not active_frames:
        return VADProxyMetrics(
            threshold_dbfs=threshold_dbfs,
            prefix_padding_ms=prefix_padding_ms,
            min_speech_ms=min_speech_ms,
            first_voice_ms=-1,
            last_voice_ms=-1,
            speech_ms=0,
            padded_start_ms=-1,
            unrecovered_prefix_ms=-1,
            initial_200ms_dbfs=initial_200ms_dbfs,
            initial_500ms_dbfs=initial_500ms_dbfs,
            low_start_risk=True,
        )

    first_voice_ms = active_frames[0] * frame_ms
    last_voice_ms = (active_frames[-1] + 1) * frame_ms
    speech_ms = max(0, last_voice_ms - first_voice_ms)
    padded_start_ms = max(0, first_voice_ms - prefix_padding_ms)
    unrecovered_prefix_ms = padded_start_ms
    low_start_risk = (
        speech_ms < min_speech_ms
        or first_voice_ms > 0
        or unrecovered_prefix_ms > 0
        or initial_200ms_dbfs < threshold_dbfs
    )
    return VADProxyMetrics(
        threshold_dbfs=threshold_dbfs,
        prefix_padding_ms=prefix_padding_ms,
        min_speech_ms=min_speech_ms,
        first_voice_ms=first_voice_ms,
        last_voice_ms=last_voice_ms,
        speech_ms=speech_ms,
        padded_start_ms=padded_start_ms,
        unrecovered_prefix_ms=unrecovered_prefix_ms,
        initial_200ms_dbfs=initial_200ms_dbfs,
        initial_500ms_dbfs=initial_500ms_dbfs,
        low_start_risk=low_start_risk,
    )


def read_wav_mono16(path: Path | str, *, target_sample_rate: int = TARGET_SAMPLE_RATE) -> AudioSample:
    wav_path = Path(path)
    with wave.open(str(wav_path), "rb") as handle:
        channels = handle.getnchannels()
        sample_width = handle.getsampwidth()
        sample_rate = handle.getframerate()
        pcm = handle.readframes(handle.getnframes())

    if sample_width != SAMPLE_WIDTH:
        raise ValueError(f"{wav_path} must be PCM16 WAV, got sample_width={sample_width}")
    if channels != TARGET_CHANNELS:
        pcm = audioop.tomono(pcm, SAMPLE_WIDTH, 0.5, 0.5)
        channels = TARGET_CHANNELS
    if sample_rate != target_sample_rate:
        pcm, _ = audioop.ratecv(pcm, SAMPLE_WIDTH, channels, sample_rate, target_sample_rate, None)
        sample_rate = target_sample_rate

    return AudioSample(name=wav_path.stem, pcm=pcm, sample_rate=sample_rate, channels=channels)


def write_wav(path: Path | str, sample: AudioSample | bytes, *, sample_rate: int = TARGET_SAMPLE_RATE) -> Path:
    wav_path = Path(path)
    wav_path.parent.mkdir(parents=True, exist_ok=True)
    if isinstance(sample, AudioSample):
        pcm = sample.pcm
        sample_rate = sample.sample_rate
        channels = sample.channels
    else:
        pcm = sample
        channels = TARGET_CHANNELS

    with wave.open(str(wav_path), "wb") as handle:
        handle.setnchannels(channels)
        handle.setsampwidth(SAMPLE_WIDTH)
        handle.setframerate(sample_rate)
        handle.writeframes(pcm)
    return wav_path


def pcm_to_wav_bytes(sample: AudioSample) -> bytes:
    import io

    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as handle:
        handle.setnchannels(sample.channels)
        handle.setsampwidth(SAMPLE_WIDTH)
        handle.setframerate(sample.sample_rate)
        handle.writeframes(sample.pcm)
    return buffer.getvalue()


def _mulaw_roundtrip(pcm: bytes) -> bytes:
    mulaw = audioop.lin2ulaw(pcm, SAMPLE_WIDTH)
    return audioop.ulaw2lin(mulaw, SAMPLE_WIDTH)


def _apply_gain(pcm: bytes, gain: float) -> bytes:
    return audioop.mul(pcm, SAMPLE_WIDTH, gain)


def _add_gaussian_noise_to_pcm(pcm: bytes, *, snr_db: float, seed: int) -> bytes:
    samples = _samples_from_pcm(pcm)
    if not samples:
        return pcm

    rng = random.Random(seed)
    signal_rms = math.sqrt(sum(float(item) * item for item in samples) / len(samples))
    noise_rms = signal_rms / (10 ** (snr_db / 20.0)) if signal_rms > 0 else 0.0
    out = []
    for value in samples:
        out.append(value + rng.gauss(0.0, noise_rms))
    return _pcm_from_samples(out)


def _noise_pcm(*, duration_ms: int, sample_rate: int, dbfs: float, seed: int) -> bytes:
    rng = random.Random(seed)
    amplitude = 32768 * (10 ** (dbfs / 20.0))
    samples = round(sample_rate * duration_ms / 1000)
    return _pcm_from_samples([rng.gauss(0.0, amplitude) for _ in range(samples)])


def _add_silence(sample: AudioSample, *, before_ms: int, after_ms: int) -> AudioSample:
    before = _silence_pcm(sample.sample_rate, before_ms)
    after = _silence_pcm(sample.sample_rate, after_ms)
    return AudioSample(name="leading_trailing_silence", pcm=before + sample.pcm + after)


def _add_noise(sample: AudioSample, *, snr_db: float, seed: int) -> AudioSample:
    return AudioSample(
        name=f"noise_snr_{round(snr_db)}",
        pcm=_add_gaussian_noise_to_pcm(sample.pcm, snr_db=snr_db, seed=seed),
    )


def _with_name(name: str, pcm: bytes, sample: AudioSample) -> AudioSample:
    return AudioSample(name=name, pcm=pcm, sample_rate=sample.sample_rate, channels=sample.channels)


def _add_named_silence(sample: AudioSample, *, name: str, before_ms: int, after_ms: int) -> AudioSample:
    return _with_name(
        name,
        _silence_pcm(sample.sample_rate, before_ms) + sample.pcm + _silence_pcm(sample.sample_rate, after_ms),
        sample,
    )


def _add_prefix_noise(sample: AudioSample, *, name: str, before_ms: int, dbfs: float, seed: int) -> AudioSample:
    return _with_name(
        name,
        _noise_pcm(duration_ms=before_ms, sample_rate=sample.sample_rate, dbfs=dbfs, seed=seed) + sample.pcm,
        sample,
    )


def _apply_prefix_gain(sample: AudioSample, *, name: str, prefix_ms: int, gain: float) -> AudioSample:
    prefix, rest = _split_pcm_at_ms(sample.pcm, sample_rate=sample.sample_rate, split_ms=prefix_ms)
    return _with_name(name, _apply_gain(prefix, gain) + rest, sample)


def _fade_in(sample: AudioSample, *, name: str, duration_ms: int) -> AudioSample:
    samples = _samples_from_pcm(sample.pcm)
    fade_samples = max(1, round(sample.sample_rate * duration_ms / 1000))
    out = []
    for idx, value in enumerate(samples):
        if idx < fade_samples:
            gain = idx / fade_samples
            out.append(value * gain)
        else:
            out.append(value)
    return _with_name(name, _pcm_from_samples(out), sample)


def _prefix_gain_with_noise(
    sample: AudioSample,
    *,
    name: str,
    prefix_ms: int,
    gain: float,
    snr_db: float,
    seed: int,
) -> AudioSample:
    prefix, rest = _split_pcm_at_ms(sample.pcm, sample_rate=sample.sample_rate, split_ms=prefix_ms)
    quiet_prefix = _apply_gain(prefix, gain)
    noisy_prefix = _add_gaussian_noise_to_pcm(quiet_prefix, snr_db=snr_db, seed=seed)
    return _with_name(name, noisy_prefix + rest, sample)


def build_variations(sample: AudioSample) -> list[AudioSample]:
    return [
        AudioSample(name="clean", pcm=sample.pcm),
        AudioSample(name="low_volume", pcm=_apply_gain(sample.pcm, 0.45)),
        AudioSample(name="very_low_volume", pcm=_apply_gain(sample.pcm, 0.22)),
        AudioSample(name="high_volume", pcm=_apply_gain(sample.pcm, 1.55)),
        AudioSample(name="clipped_high_volume", pcm=_apply_gain(sample.pcm, 2.6)),
        _add_silence(sample, before_ms=600, after_ms=900),
        _add_named_silence(sample, name="short_leading_silence", before_ms=120, after_ms=400),
        _add_named_silence(sample, name="long_leading_silence", before_ms=1600, after_ms=900),
        _add_noise(sample, snr_db=20, seed=20),
        _add_noise(sample, snr_db=15, seed=15),
        _add_noise(sample, snr_db=10, seed=10),
        _add_prefix_noise(sample, name="pre_noise_300ms", before_ms=300, dbfs=-34.0, seed=300),
        _add_prefix_noise(sample, name="pre_noise_800ms", before_ms=800, dbfs=-34.0, seed=800),
        AudioSample(name="telephony_profile", pcm=_mulaw_roundtrip(sample.pcm)),
        AudioSample(name="telephony_low_volume", pcm=_apply_gain(_mulaw_roundtrip(sample.pcm), 0.45)),
        _fade_in(sample, name="initial_fade_in_250ms", duration_ms=250),
        _fade_in(sample, name="initial_fade_in_500ms", duration_ms=500),
        _fade_in(sample, name="initial_fade_in_900ms", duration_ms=900),
        _apply_prefix_gain(sample, name="initial_dip_300ms", prefix_ms=300, gain=0.20),
        _apply_prefix_gain(sample, name="initial_dip_700ms", prefix_ms=700, gain=0.25),
        _apply_prefix_gain(sample, name="prefix_300ms_15pct", prefix_ms=300, gain=0.15),
        _apply_prefix_gain(sample, name="prefix_600ms_20pct", prefix_ms=600, gain=0.20),
        _apply_prefix_gain(sample, name="prefix_900ms_25pct", prefix_ms=900, gain=0.25),
        _prefix_gain_with_noise(
            sample,
            name="low_prefix_noise_600ms",
            prefix_ms=600,
            gain=0.22,
            snr_db=8,
            seed=608,
        ),
        _with_name(
            "low_prefix_telephony_700ms",
            _mulaw_roundtrip(_apply_prefix_gain(sample, name="unused", prefix_ms=700, gain=0.25).pcm),
            sample,
        ),
    ]
