"""Tests for autonomous AgentRunner batch execution."""

from types import SimpleNamespace
from typing import Any, AsyncIterator, Dict, List, Optional, cast
from unittest.mock import AsyncMock

import pytest
from openai.types.chat.chat_completion_message_tool_call import (
    ChatCompletionMessageToolCall,
    Function,
)

from swarm.agent import Agent
from swarm.context import RunContext
from swarm.providers.protocols import ModelRouterProtocol
from swarm.runners.agent import AgentRunner, AgentRunResult, AgentSignal
from swarm.tool import ResultMetadata, ToolExecutionResponse


class _FakeContext:
    def __init__(
        self,
        *,
        conversation_id: str = "conv-1",
        context_variables: Optional[Dict[str, Any]] = None,
        history: Optional[List[Any]] = None,
    ) -> None:
        self._conversation_id = conversation_id
        self._context_variables = context_variables or {}
        self._history = history or []
        self.workflow_context: Dict[str, Any] = {}
        self.user_id = "user-1"
        self.audience = "consumer"
        self.node_key = "agent_node"
        self._designer_fullname = "Agent Runner"

    @property
    def conversation_id(self) -> str:
        return self._conversation_id

    @property
    def current_context_variables(self) -> Dict[str, Any]:
        return self._context_variables

    def get_history_for_completion(self) -> List[Any]:
        return self._history

    @property
    def designer_fullname(self) -> str:
        return self._designer_fullname


class _FakeStreamManager:
    def __init__(self) -> None:
        self.statuses: List[str] = []

    async def publish_status(self, status: str) -> None:
        self.statuses.append(status)

    async def publish_error(self, message: str, code: int = 500) -> None:
        _ = (message, code)

    async def publish_message(self, content: Any, **kwargs: Any) -> None:
        _ = (content, kwargs)


class _FakeMessageStore:
    def __init__(self) -> None:
        self.calls: List[Dict[str, Any]] = []

    async def persist_messages(
        self,
        *,
        messages: Any,
        conversation_id: str,
        context_variables: Dict[str, Any],
        designer_model: str,
        designer_fullname: str,
        user_id: Optional[str] = None,
        audience: Optional[str] = None,
        node_key: Optional[str] = None,
        is_client_message_override: Optional[bool] = None,
    ) -> None:
        self.calls.append(
            {
                "messages": messages,
                "conversation_id": conversation_id,
                "context_variables": context_variables,
                "designer_model": designer_model,
                "designer_fullname": designer_fullname,
                "user_id": user_id,
                "audience": audience,
                "node_key": node_key,
                "is_client_message_override": is_client_message_override,
            }
        )


class _FakeRouter:
    def __init__(self, responses: List[Any]) -> None:
        self._responses = responses.copy()
        self.requests: List[Any] = []

    async def create_completion_async(self, request: Any) -> Any:
        self.requests.append(request)
        if not self._responses:
            raise AssertionError("No fake responses remaining")
        return self._responses.pop(0)

    async def stream_completion(self, request: Any) -> AsyncIterator[Any]:
        _ = request
        raise AssertionError(
            "stream_completion should not be called in AgentRunner tests"
        )
        yield None  # pragma: no cover


def _make_tool_call(
    *,
    call_id: str = "call-1",
    name: str = "think",
    arguments: str = "{}",
) -> ChatCompletionMessageToolCall:
    return ChatCompletionMessageToolCall(
        id=call_id,
        type="function",
        function=Function(name=name, arguments=arguments),
    )


def _make_completion(
    *,
    content: Optional[str] = "done",
    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None,
    model: str = "gpt-4o",
    prompt_tokens: int = 0,
    completion_tokens: int = 0,
) -> Any:
    usage = (
        None
        if prompt_tokens == 0 and completion_tokens == 0
        else SimpleNamespace(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
        )
    )
    return cast(
        Any,
        SimpleNamespace(
            choices=[
                SimpleNamespace(
                    message=SimpleNamespace(content=content, tool_calls=tool_calls)
                )
            ],
            usage=usage,
            model=model,
        ),
    )


def _make_run_context(
    *,
    router: ModelRouterProtocol,
    context: Optional[_FakeContext] = None,
    message_store: Optional[_FakeMessageStore] = None,
) -> RunContext[_FakeContext]:
    return RunContext(
        deps=context or _FakeContext(),
        stream_manager=_FakeStreamManager(),
        message_store=message_store or _FakeMessageStore(),
        router=router,
        retry_policy=None,
    )


class TestAgentRunner:
    @pytest.mark.asyncio
    async def test_run_batch_completes_on_no_tool_calls(self) -> None:
        router = _FakeRouter(
            [
                _make_completion(
                    content="All done",
                    tool_calls=None,
                )
            ]
        )
        run_context = _make_run_context(router=router)
        execute_tools_mock = AsyncMock()

        runner = AgentRunner(execute_tools_fn=cast(Any, execute_tools_mock))
        result = await runner.run_batch(
            agent=Agent(name="Agent", model="gpt-4o", instructions="system", tools=[]),
            run_context=cast(Any, run_context),
            max_turns=3,
        )

        assert result == AgentRunResult(
            signal=AgentSignal.COMPLETED,
            turns_executed=1,
        )
        execute_tools_mock.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_run_batch_maps_completed_signal(self) -> None:
        router = _FakeRouter(
            [
                _make_completion(
                    content=None,
                    tool_calls=[_make_tool_call(name="agent_completed")],
                )
            ]
        )
        run_context = _make_run_context(router=router)
        execute_tools_mock = AsyncMock(
            return_value=ToolExecutionResponse(
                messages=[],
                context_variables={},
                metadata=ResultMetadata(
                    agent_completed=True,
                    completion_summary="Finished work",
                ),
            )
        )

        runner = AgentRunner(execute_tools_fn=cast(Any, execute_tools_mock))
        result = await runner.run_batch(
            agent=Agent(name="Agent", model="gpt-4o", instructions="system", tools=[]),
            run_context=cast(Any, run_context),
            max_turns=2,
        )

        assert result.signal == AgentSignal.COMPLETED
        assert result.turns_executed == 1
        assert result.completion_summary == "Finished work"
        execute_tools_mock.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_run_batch_maps_paused_signal(self) -> None:
        router = _FakeRouter(
            [
                _make_completion(
                    content=None,
                    tool_calls=[_make_tool_call(name="ask_user")],
                )
            ]
        )
        run_context = _make_run_context(router=router)
        execute_tools_mock = AsyncMock(
            return_value=ToolExecutionResponse(
                messages=[],
                context_variables={},
                metadata=ResultMetadata(
                    agent_paused=True,
                    pause_reason="Need confirmation",
                ),
            )
        )

        runner = AgentRunner(execute_tools_fn=cast(Any, execute_tools_mock))
        result = await runner.run_batch(
            agent=Agent(name="Agent", model="gpt-4o", instructions="system", tools=[]),
            run_context=cast(Any, run_context),
            max_turns=2,
        )

        assert result.signal == AgentSignal.PAUSED
        assert result.turns_executed == 1
        assert result.pause_reason == "Need confirmation"

    @pytest.mark.asyncio
    async def test_run_batch_returns_max_turns_when_budget_exhausted(self) -> None:
        router = _FakeRouter(
            [
                _make_completion(
                    content=None, tool_calls=[_make_tool_call(call_id="c1")]
                ),
                _make_completion(
                    content=None, tool_calls=[_make_tool_call(call_id="c2")]
                ),
            ]
        )
        run_context = _make_run_context(router=router)
        execute_tools_mock = AsyncMock(
            return_value=ToolExecutionResponse(
                messages=[],
                context_variables={},
                metadata=ResultMetadata(),
            )
        )

        runner = AgentRunner(execute_tools_fn=cast(Any, execute_tools_mock))
        result = await runner.run_batch(
            agent=Agent(name="Agent", model="gpt-4o", instructions="system", tools=[]),
            run_context=cast(Any, run_context),
            max_turns=2,
        )

        assert result.signal == AgentSignal.MAX_TURNS
        assert result.turns_executed == 2
        assert execute_tools_mock.await_count == 2

    @pytest.mark.asyncio
    async def test_run_batch_attaches_usage_to_persisted_assistant_message(
        self,
    ) -> None:
        router = _FakeRouter(
            [
                _make_completion(
                    content="Done",
                    tool_calls=None,
                    model="gpt-4o-2024-08-06",
                    prompt_tokens=120,
                    completion_tokens=45,
                )
            ]
        )
        message_store = _FakeMessageStore()
        run_context = _make_run_context(router=router, message_store=message_store)

        runner = AgentRunner(execute_tools_fn=cast(Any, AsyncMock()))
        result = await runner.run_batch(
            agent=Agent(name="Agent", model="gpt-4o", instructions="system", tools=[]),
            run_context=cast(Any, run_context),
            max_turns=1,
        )

        assert result.signal == AgentSignal.COMPLETED
        assert len(message_store.calls) == 1
        persisted_messages = message_store.calls[0]["messages"]
        assert len(persisted_messages) == 1
        assistant_message = persisted_messages[0]
        assert assistant_message.usage is not None
        assert assistant_message.usage.model == "gpt-4o-2024-08-06"
        assert assistant_message.usage.input_tokens == 120
        assert assistant_message.usage.output_tokens == 45
