"""Integration tests for RetryMiddleware with the Agent execution loop."""

from __future__ import annotations

from collections.abc import AsyncGenerator
from typing import Any, Dict, List, Optional
from unittest.mock import AsyncMock, patch

import pytest

from agentbyte.agents import Agent, AgentResponse
from agentbyte.agents.types import ToolApprovalEvent, 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 RetryMiddleware
from agentbyte.tools import ApprovalMode, BaseTool
from agentbyte.types import ToolResult


# ---------------------------------------------------------------------------
# Shared fakes (mirrors test_agent_basic.py patterns)
# ---------------------------------------------------------------------------


class FakeClientConfig(BaseChatCompletionClientConfig):
    pass


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

    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]:
        yield ChatCompletionChunk(content="done", 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=[])


class FlakyTool(BaseTool):
    """Non-streaming tool that raises for the first *fail_times* calls then succeeds."""

    def __init__(self, fail_times: int = 1, result: str = "42") -> None:
        super().__init__(name="flaky_tool", description="Fails then succeeds")
        self._fail_times = fail_times
        self._call_count = 0
        self.result = result

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

    async def execute(self, parameters: Dict[str, Any]) -> ToolResult:
        self._call_count += 1
        if self._call_count <= self._fail_times:
            raise RuntimeError(f"transient failure #{self._call_count}")
        return ToolResult(success=True, result=self.result)


class StreamingRetryBeforeOutputTool(BaseTool):
    """Streaming tool that fails before emitting output on the first call."""

    def __init__(self) -> None:
        super().__init__(name="streaming_flaky", description="Streaming tool")
        self._call_count = 0

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

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

    async def execute_stream(
        self,
        parameters: Dict[str, Any],
        cancellation_token: Optional[Any] = None,
    ) -> AsyncGenerator[ToolResult | Dict[str, Any], None]:
        self._call_count += 1
        if self._call_count == 1:
            raise RuntimeError("stream failed before output")

        yield {"progress": "started"}
        yield ToolResult(success=True, result="streamed_ok")


class StreamingFailAfterOutputTool(BaseTool):
    """Streaming tool that emits output, then fails, and must not be retried."""

    def __init__(self) -> None:
        super().__init__(name="streaming_flaky", description="Streaming tool")
        self._call_count = 0

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

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

    async def execute_stream(
        self,
        parameters: Dict[str, Any],
        cancellation_token: Optional[Any] = None,
    ) -> AsyncGenerator[ToolResult | Dict[str, Any], None]:
        self._call_count += 1
        yield {"progress": "started"}
        raise RuntimeError("stream failed after output")


class ApprovalStreamingTool(BaseTool):
    """Streaming tool that requires approval and should never execute before approval."""

    def __init__(self) -> None:
        super().__init__(
            name="streaming_flaky",
            description="Streaming tool",
            approval_mode=ApprovalMode.ALWAYS,
        )
        self._call_count = 0

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

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

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


# ---------------------------------------------------------------------------
# Helper: build standard LLM responses for one tool call + final answer
# ---------------------------------------------------------------------------


def _tool_call_then_answer(tool_name: str, answer: str) -> List[ChatCompletionResult]:
    return [
        ChatCompletionResult(
            message=AssistantMessage(
                content="calling tool",
                source="fake",
                tool_calls=[
                    {"tool_name": tool_name, "parameters": {}, "call_id": "call-1"}
                ],
            ),
            usage=Usage(tokens_input=5, tokens_output=2),
            model="fake",
        ),
        ChatCompletionResult(
            message=AssistantMessage(content=answer, source="fake"),
            usage=Usage(tokens_input=5, tokens_output=2),
            model="fake",
        ),
    ]


# ---------------------------------------------------------------------------
# Test A: non-streaming tool_call retried through RetryMiddleware
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_agent_retries_non_streaming_tool_on_transient_failure():
    """RetryMiddleware retries a failing non-streaming tool until it succeeds."""
    tool = FlakyTool(fail_times=2, result="42")
    client = FakeModelClient(_tool_call_then_answer("flaky_tool", "the answer is 42"))

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[tool],
        middlewares=[RetryMiddleware(max_retries=3, initial_interval=0.1, jitter=False)],
    )

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        response = await agent.run(UserMessage(content="run flaky tool", source="user"))

    assert response.finish_reason == "stop"
    assert response.messages[-1].content == "the answer is 42"
    # Tool was called 3 times: 2 failures + 1 success
    assert tool._call_count == 3
    # asyncio.sleep called twice (after failure 1 and failure 2)
    assert mock_sleep.call_count == 2


@pytest.mark.asyncio
async def test_agent_surfaces_tool_error_after_retries_exhausted():
    """After max_retries the tool error is surfaced in the conversation as a ToolMessage."""
    tool = FlakyTool(fail_times=99, result="never")
    client = FakeModelClient(_tool_call_then_answer("flaky_tool", "could not complete"))

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[tool],
        middlewares=[RetryMiddleware(max_retries=2, initial_interval=0.1, jitter=False)],
    )

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        response = await agent.run(UserMessage(content="run flaky tool", source="user"))

    # Tool was called max_retries times (2)
    assert tool._call_count == 2
    assert mock_sleep.call_count == 1
    # A ToolMessage with the error should be in the conversation
    tool_messages = [m for m in response.messages if m.role == "tool"]
    assert len(tool_messages) == 1
    assert tool_messages[0].success is False


# ---------------------------------------------------------------------------
# Test B: streaming tool execution now uses the shared retry-aware chain
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_streaming_tool_retries_before_first_emitted_item():
    """Streaming tools are retried when the failure happens before output begins."""
    tool = StreamingRetryBeforeOutputTool()
    client = FakeModelClient(_tool_call_then_answer("streaming_flaky", "done"))

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[tool],
        middlewares=[RetryMiddleware(max_retries=3, jitter=False)],
    )

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        response: Optional[AgentResponse] = None
        streamed_items: list[object] = []
        response_event: Optional[ToolCallResponseEvent] = None
        async for item in agent.run_stream(
            UserMessage(content="run streaming tool", source="user"), verbose=True
        ):
            if isinstance(item, AgentResponse):
                response = item
            elif isinstance(item, ToolCallResponseEvent):
                response_event = item
            else:
                streamed_items.append(item)

    assert response is not None
    assert response.finish_reason == "stop"
    assert tool._call_count == 2
    assert mock_sleep.call_count == 1
    assert {"progress": "started"} in streamed_items
    assert response_event is not None
    assert response_event.result is not None
    assert response_event.result.result == "streamed_ok"


@pytest.mark.asyncio
async def test_streaming_tool_not_retried_after_output_has_started():
    """Streaming tools are not retried after user-visible output has been emitted."""
    tool = StreamingFailAfterOutputTool()
    client = FakeModelClient(_tool_call_then_answer("streaming_flaky", "done"))

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[tool],
        middlewares=[RetryMiddleware(max_retries=3, jitter=False)],
    )

    with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
        response = await agent.run(
            UserMessage(content="run streaming tool", source="user")
        )

    assert response.finish_reason == "stop"
    assert tool._call_count == 1
    mock_sleep.assert_not_called()
    tool_messages = [m for m in response.messages if m.role == "tool"]
    assert len(tool_messages) == 1
    assert tool_messages[0].success is False
    assert tool_messages[0].error == "stream failed after output"


@pytest.mark.asyncio
async def test_streaming_tool_with_approval_pauses_before_execution():
    """Approval-required streaming tools should still pause before any execution starts."""
    tool = ApprovalStreamingTool()
    client = FakeModelClient(_tool_call_then_answer("streaming_flaky", "done"))

    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[tool],
        middlewares=[RetryMiddleware(max_retries=3, jitter=False)],
    )

    approval_event: Optional[ToolApprovalEvent] = None
    async for item in agent.run_stream(
        UserMessage(content="run streaming tool", source="user"), verbose=True
    ):
        if isinstance(item, ToolApprovalEvent):
            approval_event = item
            break

    assert approval_event is not None
    assert tool._call_count == 0
