"""Basic Agent execution loop tests."""

from collections.abc import AsyncGenerator
from typing import Any, Dict, List, Optional

import pytest

from agentbyte.agents import Agent, AgentResponse
from agentbyte.cancellation_token import CancellationToken
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.messages import AssistantMessage, Message, Usage, UserMessage
from agentbyte.tools import BaseTool
from agentbyte.types import ToolResult


class FakeClientConfig(BaseChatCompletionClientConfig):
    pass


class FakeModelClient(BaseChatCompletionClient):
    def __init__(
        self,
        responses: List[ChatCompletionResult],
        stream_chunks: Optional[List[ChatCompletionChunk]] = None,
    ) -> None:
        super().__init__(model="fake", client=object())
        self._responses = responses
        self._index = 0
        self._stream_chunks = stream_chunks

    async def create(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        response = self._responses[self._index]
        self._index += 1
        return response

    async def create_stream(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[ChatCompletionChunk, None]:
        if self._stream_chunks is not None:
            for chunk in self._stream_chunks:
                yield chunk
            return

        yield ChatCompletionChunk(content="hello", model="fake", is_complete=False)
        yield ChatCompletionChunk(
            content="",
            model="fake",
            is_complete=True,
            usage=Usage(tokens_input=1, tokens_output=1),
        )

    def _to_config(self) -> BaseChatCompletionClientConfig:
        return FakeClientConfig(model=self.model, config=self.config)

    @classmethod
    def _from_config(cls, config: BaseChatCompletionClientConfig) -> "FakeModelClient":
        return cls(responses=[])


def add(a: int, b: int) -> int:
    return a + b


class StreamingAddTool(BaseTool):
    def __init__(self) -> None:
        super().__init__(name="stream_add", description="Streamed addition tool")

    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "a": {"type": "integer"},
                "b": {"type": "integer"},
            },
            "required": ["a", "b"],
        }

    async def execute(self, parameters: Dict[str, Any]) -> ToolResult:
        return ToolResult(success=True, result=parameters["a"] + parameters["b"])

    async def execute_stream(
        self,
        parameters: Dict[str, Any],
        cancellation_token: Optional[Any] = None,
    ) -> AsyncGenerator[ToolResult | Dict[str, Any], None]:
        yield {"progress": "started"}
        yield ToolResult(success=True, result=parameters["a"] + parameters["b"])


@pytest.mark.asyncio
async def test_agent_run_without_tools() -> None:
    result = ChatCompletionResult(
        message=AssistantMessage(content="done", source="fake"),
        usage=Usage(tokens_input=5, tokens_output=3),
        model="fake",
    )
    client = FakeModelClient([result])
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
    )

    response = await agent.run(UserMessage(content="hello", source="user"))

    assert response.finish_reason == "stop"
    assert response.messages[-1].content == "done"


@pytest.mark.asyncio
async def test_agent_run_accepts_string_task() -> None:
    result = ChatCompletionResult(
        message=AssistantMessage(content="done", source="fake"),
        usage=Usage(tokens_input=5, tokens_output=3),
        model="fake",
    )
    client = FakeModelClient([result])
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
    )

    response = await agent.run("hello")

    assert response.finish_reason == "stop"
    assert response.messages[0].role == "user"
    assert response.messages[0].source == "user"
    assert response.messages[0].content == "hello"


@pytest.mark.asyncio
async def test_agent_run_with_tool_call() -> None:
    first = ChatCompletionResult(
        message=AssistantMessage(
            content="calling tool",
            source="fake",
            tool_calls=[
                {
                    "tool_name": "add",
                    "parameters": {"a": 2, "b": 3},
                    "call_id": "call-1",
                }
            ],
        ),
        usage=Usage(tokens_input=10, tokens_output=4),
        model="fake",
    )
    second = ChatCompletionResult(
        message=AssistantMessage(content="result is 5", source="fake"),
        usage=Usage(tokens_input=8, tokens_output=2),
        model="fake",
    )

    client = FakeModelClient([first, second])
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[add],
    )

    response = await agent.run(UserMessage(content="calculate", source="user"))

    assert response.messages[-1].content == "result is 5"
    assert any(message.role == "tool" for message in response.messages)


@pytest.mark.asyncio
async def test_agent_cancelled_execution() -> None:
    result = ChatCompletionResult(
        message=AssistantMessage(content="done", source="fake"),
        usage=Usage(tokens_input=5, tokens_output=3),
        model="fake",
    )
    client = FakeModelClient([result])
    token = CancellationToken()
    token.cancel()

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
    )

    response = await agent.run(
        UserMessage(content="hello", source="user"), cancellation_token=token
    )

    assert response.finish_reason == "cancelled"


@pytest.mark.asyncio
async def test_agent_run_non_stream_aggregates_cost_estimate() -> None:
    first = ChatCompletionResult(
        message=AssistantMessage(
            content="calling tool",
            source="fake",
            tool_calls=[
                {
                    "tool_name": "add",
                    "parameters": {"a": 1, "b": 2},
                    "call_id": "call-1",
                }
            ],
        ),
        usage=Usage(tokens_input=6, tokens_output=3, cost_estimate=0.0011),
        model="fake",
    )
    second = ChatCompletionResult(
        message=AssistantMessage(content="result is 3", source="fake"),
        usage=Usage(tokens_input=4, tokens_output=2, cost_estimate=0.0022),
        model="fake",
    )

    client = FakeModelClient([first, second])
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[add],
    )

    response = await agent.run(UserMessage(content="calculate", source="user"))

    assert response.finish_reason == "stop"
    assert response.usage.llm_calls == 2
    assert response.usage.cost_estimate == pytest.approx(0.0033)


@pytest.mark.asyncio
async def test_agent_run_stream_aggregates_cost_estimate() -> None:
    stream_chunks = [
        ChatCompletionChunk(content="hel", model="fake", is_complete=False),
        ChatCompletionChunk(content="lo", model="fake", is_complete=False),
        ChatCompletionChunk(
            content="",
            model="fake",
            is_complete=True,
            usage=Usage(tokens_input=3, tokens_output=2, cost_estimate=0.0042),
        ),
    ]

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=FakeModelClient(responses=[], stream_chunks=stream_chunks),
    )

    final_response: Optional[AgentResponse] = None
    async for item in agent.run_stream(
        UserMessage(content="hi", source="user"), stream_tokens=True
    ):
        if isinstance(item, AgentResponse):
            final_response = item

    assert final_response is not None
    assert final_response.finish_reason == "stop"
    assert final_response.usage.llm_calls == 1
    assert final_response.usage.cost_estimate == pytest.approx(0.0042)


@pytest.mark.asyncio
async def test_agent_run_stream_accepts_string_task() -> None:
    stream_chunks = [
        ChatCompletionChunk(content="hel", model="fake", is_complete=False),
        ChatCompletionChunk(content="lo", model="fake", is_complete=False),
        ChatCompletionChunk(
            content="",
            model="fake",
            is_complete=True,
            usage=Usage(tokens_input=3, tokens_output=2),
        ),
    ]

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=FakeModelClient(responses=[], stream_chunks=stream_chunks),
    )

    first_user_message: Optional[Message] = None
    final_response: Optional[AgentResponse] = None
    async for item in agent.run_stream("hi", stream_tokens=True):
        if isinstance(item, Message) and item.role == "user" and first_user_message is None:
            first_user_message = item
        if isinstance(item, AgentResponse):
            final_response = item

    assert first_user_message is not None
    assert first_user_message.source == "user"
    assert first_user_message.content == "hi"
    assert final_response is not None
    assert final_response.finish_reason == "stop"


@pytest.mark.asyncio
async def test_agent_run_with_streaming_tool_path() -> None:
    first = ChatCompletionResult(
        message=AssistantMessage(
            content="calling streamed tool",
            source="fake",
            tool_calls=[
                {
                    "tool_name": "stream_add",
                    "parameters": {"a": 2, "b": 3},
                    "call_id": "call-1",
                }
            ],
        ),
        usage=Usage(tokens_input=5, tokens_output=2),
        model="fake",
    )
    second = ChatCompletionResult(
        message=AssistantMessage(content="done", source="fake"),
        usage=Usage(tokens_input=4, tokens_output=1),
        model="fake",
    )

    client = FakeModelClient([first, second])
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[StreamingAddTool()],
    )

    seen_progress = False
    response: Optional[AgentResponse] = None
    async for item in agent.run_stream(
        UserMessage(content="calculate", source="user"), verbose=False
    ):
        if isinstance(item, dict) and item.get("progress") == "started":
            seen_progress = True
        if isinstance(item, AgentResponse):
            response = item

    assert seen_progress is True
    assert response is not None
    assert response.finish_reason == "stop"
    tool_messages = [message for message in response.messages if message.role == "tool"]
    assert len(tool_messages) == 1
    assert tool_messages[0].content == "5"
