from __future__ import annotations

import asyncio
import wave
from functools import lru_cache
from pathlib import Path
from typing import Any, AsyncIterator

DEFAULT_TAIL_SILENCE_MS = 320


@lru_cache(maxsize=16)
def _read_wav_pcm(path: str) -> tuple[bytes, int, int]:
    wav_path = Path(path)
    with wave.open(str(wav_path), "rb") as wav:
        channels = wav.getnchannels()
        sample_width = wav.getsampwidth()
        sample_rate = wav.getframerate()
        if channels != 1 or sample_width != 2:
            raise ValueError(
                f"Unsupported wait audio format for {wav_path}: "
                f"expected PCM16 mono, got channels={channels}, sample_width={sample_width}"
            )
        data = wav.readframes(10**9)
    if not data:
        raise ValueError(f"Wait audio file is empty: {wav_path}")
    return data, sample_rate, channels


async def wav_audio_frames(
    path: str,
    *,
    frame_duration_ms: int = 20,
    tail_silence_ms: int = DEFAULT_TAIL_SILENCE_MS,
) -> AsyncIterator[Any]:
    from livekit import rtc

    data, sample_rate, channels = _read_wav_pcm(path)
    samples_per_frame = max(1, round(sample_rate * max(1, frame_duration_ms) / 1000))
    bytes_per_sample = 2 * channels
    frame_bytes = samples_per_frame * bytes_per_sample

    def _frame(chunk: bytes) -> Any:
        if len(chunk) < frame_bytes:
            chunk = chunk + (b"\x00" * (frame_bytes - len(chunk)))
        return rtc.AudioFrame(
            data=chunk,
            sample_rate=sample_rate,
            num_channels=channels,
            samples_per_channel=len(chunk) // bytes_per_sample,
        )

    for offset in range(0, len(data), frame_bytes):
        chunk = data[offset : offset + frame_bytes]
        samples_per_channel = len(chunk) // bytes_per_sample
        if samples_per_channel <= 0:
            continue
        yield _frame(chunk)
        await asyncio.sleep(0)

    tail_ms = max(0, int(tail_silence_ms))
    frame_ms = max(1, int(frame_duration_ms))
    tail_frames = (tail_ms + frame_ms - 1) // frame_ms
    silence = b"\x00" * frame_bytes
    for _ in range(tail_frames):
        yield _frame(silence)
        await asyncio.sleep(0)


def wav_duration_ms(path: str) -> int:
    data, sample_rate, channels = _read_wav_pcm(path)
    samples = len(data) // (2 * channels)
    return round((samples / sample_rate) * 1000)
