"""Tests for richer verbose stream event payloads."""

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

import pytest

from agentbyte.agents import (
    Agent,
    ModelCallEvent,
    ModelResponseEvent,
    ModelStreamChunkEvent,
    TaskCompleteEvent,
    TaskStartEvent,
    ToolCallEvent,
    ToolCallResponseEvent,
)
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.messages import AssistantMessage, Message, Usage, UserMessage
from agentbyte.middleware import BaseMiddleware, MiddlewareContext


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 or []
        self.create_calls = 0
        self.create_stream_calls = 0

    async def create(
        self,
        messages: List[Message],
        tools: Optional[List[Dict[str, Any]]] = None,
        output_format: Optional[type] = None,
        **kwargs: Any,
    ) -> ChatCompletionResult:
        self.create_calls += 1
        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]:
        self.create_stream_calls += 1
        for chunk in self._stream_chunks:
            yield chunk

    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 PassThroughMiddleware(BaseMiddleware):
    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        return context

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        return result

    async def process_error(
        self, context: MiddlewareContext, error: Exception
    ) -> Optional[Any]:
        return None


@pytest.mark.asyncio
async def test_verbose_non_stream_events_include_rich_payload() -> None:
    first = ChatCompletionResult(
        message=AssistantMessage(
            content="calling add",
            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="done", source="fake"),
        usage=Usage(tokens_input=3, tokens_output=2),
        model="fake",
    )

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

    events: List[object] = []
    async for item in agent.run_stream(
        UserMessage(content="calculate", source="user"),
        verbose=True,
        stream_tokens=False,
    ):
        events.append(item)

    task_start = next(item for item in events if isinstance(item, TaskStartEvent))
    model_call = next(item for item in events if isinstance(item, ModelCallEvent))
    model_response = next(
        item for item in events if isinstance(item, ModelResponseEvent)
    )
    tool_call = next(item for item in events if isinstance(item, ToolCallEvent))
    tool_response = next(
        item for item in events if isinstance(item, ToolCallResponseEvent)
    )
    task_complete = next(item for item in events if isinstance(item, TaskCompleteEvent))

    assert task_start.task == "calculate"
    assert len(model_call.input_messages) >= 2
    assert model_call.model == "fake"
    assert model_response.has_tool_calls is True
    assert model_response.response == "calling add"
    assert tool_call.parameters == {"a": 2, "b": 3}
    assert tool_call.call_id == "call-1"
    assert tool_response.call_id == "call-1"
    assert tool_response.result is not None
    assert tool_response.result.success is True
    assert task_complete.result == "done"


@pytest.mark.asyncio
async def test_verbose_stream_events_include_chunk_events() -> 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=1, tokens_output=1),
        ),
    ]

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

    events: List[object] = []
    async for item in agent.run_stream(
        UserMessage(content="hi", source="user"), verbose=True, stream_tokens=True
    ):
        events.append(item)

    chunk_events = [item for item in events if isinstance(item, ModelStreamChunkEvent)]
    assert len(chunk_events) == 3
    assert chunk_events[0].chunk == "hel"
    assert chunk_events[0].is_final is False
    assert chunk_events[-1].is_final is True


@pytest.mark.asyncio
async def test_stream_tokens_disabled_when_middleware_configured() -> None:
    response = ChatCompletionResult(
        message=AssistantMessage(content="from-create", source="fake"),
        usage=Usage(tokens_input=2, tokens_output=1),
        model="fake",
    )
    stream_chunks = [
        ChatCompletionChunk(content="from", model="fake", is_complete=False),
        ChatCompletionChunk(content="-stream", model="fake", is_complete=False),
        ChatCompletionChunk(content="", model="fake", is_complete=True),
    ]
    client = FakeModelClient(responses=[response], stream_chunks=stream_chunks)

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        middlewares=[PassThroughMiddleware()],
    )

    events: List[object] = []
    with pytest.warns(UserWarning, match="stream_tokens=True is disabled"):
        async for item in agent.run_stream(
            UserMessage(content="hi", source="user"), verbose=True, stream_tokens=True
        ):
            events.append(item)

    chunk_events = [item for item in events if isinstance(item, ModelStreamChunkEvent)]
    assistant_messages = [item for item in events if isinstance(item, AssistantMessage)]

    assert len(chunk_events) == 0
    assert assistant_messages
    assert assistant_messages[-1].content == "from-create"
    assert client.create_calls == 1
    assert client.create_stream_calls == 0
