import asyncio
import zmq
import json
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from control_backend.agents.ri_command_agent import RICommandAgent
from control_backend.schemas.ri_message import SpeechCommand


@pytest.mark.asyncio
async def test_setup_bind(monkeypatch):
    """Test setup with bind=True"""
    fake_socket = MagicMock()
    monkeypatch.setattr(
        "control_backend.agents.ri_command_agent.context.socket", lambda _: fake_socket
    )

    agent = RICommandAgent("test@server", "password", address="tcp://localhost:5555", bind=True)
    monkeypatch.setattr(
        "control_backend.agents.ri_command_agent.settings",
        MagicMock(zmq_settings=MagicMock(internal_comm_address="tcp://internal:1234")),
    )

    await agent.setup()

    # Ensure PUB socket bound
    fake_socket.bind.assert_any_call("tcp://localhost:5555")
    # Ensure SUB socket connected to internal address and subscribed
    fake_socket.connect.assert_any_call("tcp://internal:1234")
    fake_socket.setsockopt.assert_any_call(zmq.SUBSCRIBE, b"command")

    # Ensure behaviour attached
    assert any(isinstance(b, agent.SendCommandsBehaviour) for b in agent.behaviours)


@pytest.mark.asyncio
async def test_setup_connect(monkeypatch):
    """Test setup with bind=False"""
    fake_socket = MagicMock()
    monkeypatch.setattr(
        "control_backend.agents.ri_command_agent.context.socket", lambda _: fake_socket
    )

    agent = RICommandAgent("test@server", "password", address="tcp://localhost:5555", bind=False)
    monkeypatch.setattr(
        "control_backend.agents.ri_command_agent.settings",
        MagicMock(zmq_settings=MagicMock(internal_comm_address="tcp://internal:1234")),
    )

    await agent.setup()

    # Ensure PUB socket connected
    fake_socket.connect.assert_any_call("tcp://localhost:5555")


@pytest.mark.asyncio
async def test_send_commands_behaviour_valid_message():
    """Test behaviour with valid JSON message"""
    fake_socket = AsyncMock()
    message_dict = {"message": "hello"}
    fake_socket.recv_multipart = AsyncMock(
        return_value=(b"command", json.dumps(message_dict).encode("utf-8"))
    )
    fake_socket.send_json = AsyncMock()

    agent = RICommandAgent("test@server", "password")
    agent.subsocket = fake_socket
    agent.pubsocket = fake_socket

    behaviour = agent.SendCommandsBehaviour()
    behaviour.agent = agent

    with patch("control_backend.agents.ri_command_agent.SpeechCommand") as MockSpeechCommand:
        mock_message = MagicMock()
        MockSpeechCommand.model_validate.return_value = mock_message

        await behaviour.run()

    fake_socket.recv_multipart.assert_awaited()
    fake_socket.send_json.assert_awaited_with(mock_message.model_dump())


@pytest.mark.asyncio
async def test_send_commands_behaviour_invalid_message(caplog):
    """Test behaviour with invalid JSON message triggers error logging"""
    fake_socket = AsyncMock()
    fake_socket.recv_multipart = AsyncMock(return_value=(b"command", b"{invalid_json}"))
    fake_socket.send_json = AsyncMock()

    agent = RICommandAgent("test@server", "password")
    agent.subsocket = fake_socket
    agent.pubsocket = fake_socket

    behaviour = agent.SendCommandsBehaviour()
    behaviour.agent = agent

    with caplog.at_level("ERROR"):
        await behaviour.run()

    fake_socket.recv_multipart.assert_awaited()
    fake_socket.send_json.assert_not_awaited()
    assert "Error processing message" in caplog.text