"""Agent integration tests for middleware execute_stream behavior."""

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

import pytest

from agentbyte.agents import (
    Agent,
    AgentResponse,
    ToolApprovalEvent,
    ToolValidationEvent,
)
from agentbyte.context import ToolApprovalRequest
from agentbyte.llm.base import BaseChatCompletionClient, BaseChatCompletionClientConfig
from agentbyte.llm.types import ChatCompletionChunk, ChatCompletionResult
from agentbyte.memory import ListMemory, MemoryContent
from agentbyte.messages import (
    AssistantMessage,
    Message,
    ToolCallRequest,
    Usage,
    UserMessage,
)
from agentbyte.middleware import ApprovalMiddleware, BaseMiddleware, MiddlewareContext
from agentbyte.tools import FunctionTool


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
        self.create_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]:
        if False:
            yield ChatCompletionChunk(content="", model="fake", is_complete=True)

    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 EventEmittingMiddleware(BaseMiddleware):
    async def process_request(
        self,
        context: MiddlewareContext,
    ) -> AsyncGenerator[Any, None]:
        yield ToolValidationEvent(source="mw", tool_name="noop", is_valid=True)
        yield 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


class PauseOnModelCallMiddleware(BaseMiddleware):
    async def process_request(
        self,
        context: MiddlewareContext,
    ) -> AsyncGenerator[Any, None]:
        if context.operation == "model_call":
            request = ToolApprovalRequest(
                request_id="approval_mw",
                tool_call_id="mw-call-1",
                tool_name="middleware_gate",
                parameters={"reason": "manual review"},
                original_tool_call=ToolCallRequest(
                    tool_name="middleware_gate",
                    parameters={"reason": "manual review"},
                    call_id="mw-call-1",
                ),
            )
            yield ToolApprovalEvent(source="mw", approval_request=request)
            return

        yield 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


class OperationRecorderMiddleware(BaseMiddleware):
    def __init__(self) -> None:
        self.operations: list[str] = []

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        self.operations.append(context.operation)
        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_agent_run_stream_yields_middleware_events_when_verbose() -> None:
    model_client = FakeModelClient(
        responses=[
            ChatCompletionResult(
                message=AssistantMessage(content="done", source="fake"),
                usage=Usage(tokens_input=2, tokens_output=1),
                model="fake",
            )
        ]
    )

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

    saw_middleware_event = False
    final_response: AgentResponse | None = None

    async for item in agent.run_stream(
        UserMessage(content="hello", source="user"), verbose=True, stream_tokens=False
    ):
        if isinstance(item, ToolValidationEvent):
            saw_middleware_event = True
        if isinstance(item, AgentResponse):
            final_response = item

    assert saw_middleware_event is True
    assert final_response is not None
    assert final_response.finish_reason == "stop"


@pytest.mark.asyncio
async def test_agent_run_stream_pauses_on_middleware_tool_approval_event() -> None:
    model_client = FakeModelClient(responses=[])

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

    saw_approval_event = False
    final_response: AgentResponse | None = None

    async for item in agent.run_stream(
        UserMessage(content="hello", source="user"), verbose=False, stream_tokens=False
    ):
        if isinstance(item, ToolApprovalEvent):
            saw_approval_event = True
        if isinstance(item, AgentResponse):
            final_response = item

    assert saw_approval_event is True
    assert final_response is not None
    assert final_response.finish_reason == "approval_needed"
    assert model_client.create_calls == 0


@pytest.mark.asyncio
async def test_agent_run_stream_pauses_with_approval_middleware_for_tool_call() -> None:
    model_client = FakeModelClient(
        responses=[
            ChatCompletionResult(
                message=AssistantMessage(
                    content="Need approval",
                    source="fake",
                    tool_calls=[
                        ToolCallRequest(
                            tool_name="sensitive_tool",
                            parameters={"text": "hello"},
                            call_id="tool-call-1",
                        )
                    ],
                ),
                usage=Usage(tokens_input=3, tokens_output=2),
                model="fake",
            )
        ]
    )

    tool = FunctionTool(
        lambda text: text, name="sensitive_tool", description="Sensitive"
    )
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=model_client,
        tools=[tool],
        middlewares=[ApprovalMiddleware(tool_names=["sensitive_tool"])],
    )

    saw_approval_event = False
    final_response: AgentResponse | None = None

    async for item in agent.run_stream(
        UserMessage(content="hello", source="user"), verbose=True, stream_tokens=False
    ):
        if isinstance(item, ToolApprovalEvent):
            saw_approval_event = True
        if isinstance(item, AgentResponse):
            final_response = item

    assert saw_approval_event is True
    assert final_response is not None
    assert final_response.finish_reason == "approval_needed"


@pytest.mark.asyncio
async def test_agent_routes_memory_access_through_middleware() -> None:
    model_client = FakeModelClient(
        responses=[
            ChatCompletionResult(
                message=AssistantMessage(content="done", source="fake"),
                usage=Usage(tokens_input=2, tokens_output=1),
                model="fake",
            )
        ]
    )

    memory = ListMemory()
    await memory.add(MemoryContent(content="Remember this"))

    recorder = OperationRecorderMiddleware()
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=model_client,
        memory=memory,
        middlewares=[recorder],
    )

    final_response: AgentResponse | None = None
    async for item in agent.run_stream(
        UserMessage(content="hello", source="user"), verbose=False, stream_tokens=False
    ):
        if isinstance(item, AgentResponse):
            final_response = item

    assert "memory_access" in recorder.operations
    assert final_response is not None
    assert final_response.usage.memory_operations == 1
