import argparse
from unittest.mock import MagicMock, patch

import pytest

from debgpt.backend import AbstractBackend, ZMQBackend, create_backend, stat_messages


def test_stat_messages():
    mock_llm = MagicMock()

    # Test with size attribute
    mock_res_size = MagicMock()
    mock_res_size.size = MagicMock(return_value=10)
    mock_llm.tok.apply_chat_template = MagicMock(return_value=mock_res_size)

    messages = [
        {"role": "user", "content": "hello"},
        {"role": "assistant", "content": "hi"},
    ]
    stat = stat_messages(messages, mock_llm)
    assert "num_msgs=2" in stat
    assert "ctx_size=10" in stat
    assert "latest={'role': 'assistant', 'content': 'hi'}" in stat

    # Test with input_ids dict
    mock_res_dict = {"input_ids": mock_res_size}
    mock_llm.tok.apply_chat_template = MagicMock(return_value=mock_res_dict)

    stat = stat_messages(messages, mock_llm)
    assert "ctx_size=10" in stat

    # Test with len fallback
    mock_res_list = [1, 2, 3, 4, 5, 6, 7]
    mock_llm.tok.apply_chat_template = MagicMock(return_value=mock_res_list)

    stat = stat_messages(messages, mock_llm)
    assert "ctx_size=7" in stat


def test_abstract_backend():
    mock_args = MagicMock()
    with patch("debgpt.backend.llm.create_llm") as mock_create_llm:
        mock_create_llm.return_value = MagicMock()
        backend = AbstractBackend(mock_args)

        with pytest.raises(NotImplementedError):
            backend.listen(mock_args)

        with pytest.raises(NotImplementedError):
            backend.server()


def test_create_backend_unimplemented():
    args = argparse.Namespace(
        backend_impl="unknown", host="tcp://127.0.0.1", port=12345
    )
    with pytest.raises(NotImplementedError):
        create_backend(args)


@patch("debgpt.backend.zmq.Context")
@patch("debgpt.backend.llm.create_llm")
def test_create_backend_zmq(mock_create_llm, mock_zmq_context):
    args = argparse.Namespace(backend_impl="zmq", host="tcp://127.0.0.1", port=12345)

    mock_socket = MagicMock()
    mock_zmq_context.return_value.socket.return_value = mock_socket

    backend = create_backend(args)
    assert isinstance(backend, ZMQBackend)
    mock_socket.bind.assert_called_once_with("tcp://127.0.0.1:12345")

    # Test listen
    mock_socket.recv_json.return_value = {"test": "data"}
    listen_gen = backend.listen()
    assert next(listen_gen) == {"test": "data"}

    # Test server loop logic
    backend.listen = MagicMock(
        return_value=[
            {
                "session": [{"role": "user", "content": "hi"}],
                "temperature": 0.5,
                "top_p": 0.9,
            }
        ]
    )

    backend.llm.generate_stream = MagicMock(return_value=["hello ", "world"])

    backend.server()

    # Verify socket sends
    assert mock_socket.send_json.call_count == 3  # 2 chunks + 1 done
    mock_socket.send_json.assert_any_call({"status": "chunk", "content": "hello "})
    mock_socket.send_json.assert_any_call({"status": "chunk", "content": "world"})

    done_call = mock_socket.send_json.call_args_list[-1][0][0]
    assert done_call["status"] == "done"
    assert len(done_call["session"]) == 2
    assert done_call["session"][-1]["content"] == "hello world"


@patch("debgpt.backend.zmq.Context")
@patch("debgpt.backend.llm.create_llm")
def test_create_backend_zmq_server_error(mock_create_llm, mock_zmq_context):
    args = argparse.Namespace(backend_impl="zmq", host="tcp://127.0.0.1", port=12345)

    mock_socket = MagicMock()
    mock_zmq_context.return_value.socket.return_value = mock_socket

    backend = create_backend(args)
    backend.listen = MagicMock(return_value=[[{"role": "user", "content": "hi"}]])

    # Simulate an error during generation
    backend.llm.generate_stream.side_effect = Exception("Generation failed")

    backend.server()

    mock_socket.send_json.assert_called_once_with(
        {"status": "error", "content": "Generation failed"}
    )
