from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

import pytest

from aiops.infra.agent.middlewares.chat_model_setter import ChatModelSetterMiddleware
from tests.infra.agent.middlewares.helpers import DummyModelRequest


class ContextWithChatModelId:
    def __init__(self, chat_model_id: int | None) -> None:
        self.chat_model_id = chat_model_id


@pytest.mark.asyncio
async def test_chat_model_setter_replaces_model(dummy_model_request) -> None:
    chat_model_record = {"id": 1}
    built_model = object()
    service_llm = SimpleNamespace(
        chat_model=SimpleNamespace(get=AsyncMock(return_value=chat_model_record)),
        lang_chain=SimpleNamespace(
            build_chat_model=MagicMock(return_value=built_model)
        ),
    )
    middleware = ChatModelSetterMiddleware(service_llm)
    request = dummy_model_request(
        system_prompt="prompt",
        model="original-model",
        runtime_context=ContextWithChatModelId(chat_model_id=1),
    )
    handler = AsyncMock(return_value="response")

    result = await middleware.awrap_model_call(request, handler)

    assert result == "response"
    service_llm.chat_model.get.assert_awaited_once_with(1)
    service_llm.lang_chain.build_chat_model.assert_called_once_with(chat_model_record)
    called_request = handler.await_args[0][0]
    assert isinstance(called_request, DummyModelRequest)
    assert called_request.model is built_model
    assert called_request.system_prompt == "prompt"


@pytest.mark.asyncio
async def test_chat_model_setter_requires_context_attribute(
    dummy_model_request,
) -> None:
    service_llm = SimpleNamespace(
        chat_model=SimpleNamespace(get=AsyncMock(return_value=None)),
        lang_chain=SimpleNamespace(build_chat_model=MagicMock()),
    )
    middleware = ChatModelSetterMiddleware(service_llm)
    request = dummy_model_request(
        system_prompt="prompt", model="model", runtime_context=object()
    )
    handler = AsyncMock()

    with pytest.raises(TypeError):
        await middleware.awrap_model_call(request, handler)

    handler.assert_not_awaited()


@pytest.mark.asyncio
async def test_chat_model_setter_raises_when_chat_model_missing(
    dummy_model_request,
) -> None:
    service_llm = SimpleNamespace(
        chat_model=SimpleNamespace(get=AsyncMock(return_value=None)),
        lang_chain=SimpleNamespace(build_chat_model=MagicMock()),
    )
    middleware = ChatModelSetterMiddleware(service_llm)
    request = dummy_model_request(
        system_prompt="prompt",
        model="model",
        runtime_context=ContextWithChatModelId(chat_model_id=99),
    )
    handler = AsyncMock()

    with pytest.raises(ValueError, match="99"):
        await middleware.awrap_model_call(request, handler)

    handler.assert_not_awaited()