"""
This program has been developed by students from the bachelor Computer Science at Utrecht
University within the Software Project course.
© Copyright Utrecht University (Department of Information and Computing Sciences)
"""

import os
from unittest.mock import AsyncMock, MagicMock

import pytest
import soundfile as sf
import zmq

from control_backend.agents.perception.vad_agent import VADAgent


@pytest.fixture(autouse=True)
def patch_settings():
    from control_backend.agents.perception import vad_agent

    vad_agent.settings.behaviour_settings.vad_prob_threshold = 0.5
    vad_agent.settings.behaviour_settings.vad_non_speech_patience_chunks = 3
    vad_agent.settings.behaviour_settings.vad_initial_since_speech = 0
    vad_agent.settings.vad_settings.sample_rate_hz = 16_000


@pytest.fixture(autouse=True)
def mock_torch(mocker):
    mock_torch = mocker.patch("control_backend.agents.perception.vad_agent.torch")
    mock_torch.from_numpy.side_effect = lambda arr: arr
    return mock_torch


def get_audio_chunks() -> list[bytes]:
    curr_file = os.path.realpath(__file__)
    curr_dir = os.path.dirname(curr_file)
    file = f"{curr_dir}/speech_with_pauses_16k_1c_float32.wav"

    chunk_size = 512

    chunks = []

    with sf.SoundFile(file, "r") as f:
        assert f.samplerate == 16000
        assert f.channels == 1
        assert f.subtype == "FLOAT"

        while True:
            data = f.read(chunk_size, dtype="float32")
            if len(data) != chunk_size:
                break

            chunks.append(data.tobytes())

    return chunks


@pytest.mark.asyncio
async def test_real_audio(mocker):
    """
    Test the VAD agent with only input and output mocked. Using the real model, using real audio as
    input. Ensure that it outputs some fragments with audio.
    """
    audio_chunks = get_audio_chunks()
    audio_in_socket = AsyncMock()
    audio_in_socket.recv.side_effect = audio_chunks

    mock_poller: MagicMock = mocker.patch("control_backend.agents.perception.vad_agent.azmq.Poller")
    mock_poller.return_value.poll = AsyncMock(return_value=[(audio_in_socket, zmq.POLLIN)])

    audio_out_socket = AsyncMock()

    vad_agent = VADAgent("tcp://localhost:12345", False)
    vad_agent.audio_out_socket = audio_out_socket

    # Use a fake model that marks most chunks as speech and ends with a few silences
    silence_padding = 5
    probabilities = [1.0] * len(audio_chunks) + [0.0] * silence_padding
    chunk_bytes = audio_chunks + [b"\x00" * len(audio_chunks[0])] * silence_padding
    model_item = MagicMock()
    model_item.item.side_effect = probabilities
    vad_agent.model = MagicMock(return_value=model_item)

    class DummyPoller:
        def __init__(self, data, agent):
            self.data = data
            self.agent = agent

        async def poll(self, timeout_ms=None):
            if self.data:
                return self.data.pop(0)
            self.agent._running = False
            return None

    vad_agent.audio_in_poller = DummyPoller(chunk_bytes, vad_agent)
    vad_agent._ready = AsyncMock()
    vad_agent._running = True
    vad_agent.i_since_speech = 0

    await vad_agent._streaming_loop()

    audio_out_socket.send.assert_called()
    for args in audio_out_socket.send.call_args_list:
        assert isinstance(args[0][0], bytes)
        assert len(args[0][0]) >= 512 * 4 * 3  # Should be at least 3 chunks of audio
