from __future__ import annotations

import asyncio
import sys
import types
import wave
from pathlib import Path
from unittest import mock

from app.livekit.runtime.wav_audio import wav_audio_frames, wav_duration_ms


class _AudioFrame:
    def __init__(
        self,
        *,
        data: bytes,
        sample_rate: int,
        num_channels: int,
        samples_per_channel: int,
    ) -> None:
        self.data = data
        self.sample_rate = sample_rate
        self.num_channels = num_channels
        self.samples_per_channel = samples_per_channel


def _write_wav(path: Path, *, samples: int, sample_rate: int = 1000) -> None:
    with wave.open(str(path), "wb") as wav:
        wav.setnchannels(1)
        wav.setsampwidth(2)
        wav.setframerate(sample_rate)
        wav.writeframes(b"\x01\x02" * samples)


def test_wav_audio_frames_pads_last_frame_and_adds_tail_silence(tmp_path: Path) -> None:
    wav_path = tmp_path / "audio.wav"
    _write_wav(wav_path, samples=25)

    livekit_module = types.ModuleType("livekit")
    rtc_module = types.ModuleType("livekit.rtc")
    rtc_module.AudioFrame = _AudioFrame
    livekit_module.rtc = rtc_module

    async def _collect():
        with mock.patch.dict(sys.modules, {"livekit": livekit_module, "livekit.rtc": rtc_module}):
            return [
                frame
                async for frame in wav_audio_frames(
                    str(wav_path),
                    frame_duration_ms=20,
                    tail_silence_ms=40,
                )
            ]

    frames = asyncio.run(_collect())

    assert len(frames) == 4
    assert [frame.samples_per_channel for frame in frames] == [20, 20, 20, 20]
    assert frames[1].data[:10] == b"\x01\x02" * 5
    assert frames[1].data[10:] == b"\x00" * 30
    assert frames[2].data == b"\x00" * 40
    assert frames[3].data == b"\x00" * 40
    assert wav_duration_ms(str(wav_path)) == 25