from unittest.mock import MagicMock, AsyncMock, patch

import pytest
import zmq
from spade.agent import Agent

from control_backend.agents.vad_agent import VADAgent


@pytest.fixture
def zmq_context(mocker):
    return mocker.patch("control_backend.agents.vad_agent.zmq_context")


@pytest.fixture
def streaming(mocker):
    return mocker.patch("control_backend.agents.vad_agent.Streaming")


@pytest.mark.asyncio
async def test_normal_setup(streaming):
    """
    Test that during normal setup, the VAD agent creates a Streaming behavior and creates audio sockets.
    """
    vad_agent = VADAgent("tcp://localhost:12345", False)
    vad_agent.add_behaviour = MagicMock()

    await vad_agent.setup()

    streaming.assert_called_once()
    vad_agent.add_behaviour.assert_called_once_with(streaming.return_value)
    assert vad_agent.audio_in_socket is not None
    assert vad_agent.audio_out_socket is not None


@pytest.mark.parametrize("do_bind", [True, False])
def test_in_socket_creation(zmq_context, do_bind: bool):
    """
    Test that the VAD agent creates an audio input socket, differentiating between binding and connecting.
    """
    vad_agent = VADAgent(f"tcp://{"*" if do_bind else "localhost"}:12345", do_bind)

    vad_agent._connect_audio_in_socket()

    assert vad_agent.audio_in_socket is not None

    zmq_context.socket.assert_called_once_with(zmq.SUB)
    zmq_context.socket.return_value.setsockopt_string.assert_called_once_with(zmq.SUBSCRIBE, "")

    if do_bind:
        zmq_context.socket.return_value.bind.assert_called_once_with("tcp://*:12345")
    else:
        zmq_context.socket.return_value.connect.assert_called_once_with("tcp://localhost:12345")


def test_out_socket_creation(zmq_context):
    """
    Test that the VAD agent creates an audio output socket correctly.
    """
    vad_agent = VADAgent("tcp://localhost:12345", False)

    vad_agent._connect_audio_out_socket()

    assert vad_agent.audio_out_socket is not None

    zmq_context.socket.assert_called_once_with(zmq.PUB)
    zmq_context.socket.return_value.bind_to_random_port.assert_called_once()


@pytest.mark.asyncio
async def test_out_socket_creation_failure(zmq_context):
    """
    Test setup failure when the audio output socket cannot be created.
    """
    with patch.object(Agent, "stop", new_callable=AsyncMock) as mock_super_stop:
        zmq_context.socket.return_value.bind_to_random_port.side_effect = zmq.ZMQBindError
        vad_agent = VADAgent("tcp://localhost:12345", False)

        await vad_agent.setup()

        assert vad_agent.audio_out_socket is None
        mock_super_stop.assert_called_once()


@pytest.mark.asyncio
async def test_stop(zmq_context):
    """
    Test that when the VAD agent is stopped, the sockets are closed correctly.
    """
    vad_agent = VADAgent("tcp://localhost:12345", False)

    await vad_agent.setup()
    await vad_agent.stop()

    assert zmq_context.socket.return_value.close.call_count == 2
    assert vad_agent.audio_in_socket is None
    assert vad_agent.audio_out_socket is None