"""Tests for tool approval pause/resume workflows."""

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

import pytest

from agentbyte.agents import Agent, AgentResponse, ToolApprovalEvent
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 tool


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="", 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=[])


@tool(approval_mode="always_require")
def sensitive_add(a: int, b: int) -> int:
    """Add two integers, requiring explicit approval."""
    return a + b


def _tool_call_result() -> ChatCompletionResult:
    return ChatCompletionResult(
        message=AssistantMessage(
            content="need tool",
            source="fake",
            tool_calls=[
                {
                    "tool_name": "sensitive_add",
                    "parameters": {"a": 2, "b": 3},
                    "call_id": "call-1",
                }
            ],
        ),
        usage=Usage(tokens_input=10, tokens_output=4),
        model="fake",
    )


def _final_result(content: str = "done") -> ChatCompletionResult:
    return ChatCompletionResult(
        message=AssistantMessage(content=content, source="fake"),
        usage=Usage(tokens_input=8, tokens_output=2),
        model="fake",
    )


@pytest.mark.asyncio
async def test_tool_approval_event_and_pause() -> None:
    client = FakeModelClient([_tool_call_result()])
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[sensitive_add],
    )

    events: List[Any] = []
    final_response: Optional[AgentResponse] = None

    async for item in agent.run_stream(
        UserMessage(content="calculate", source="user"), verbose=True
    ):
        events.append(item)
        if isinstance(item, AgentResponse):
            final_response = item

    approval_events = [item for item in events if isinstance(item, ToolApprovalEvent)]
    assert len(approval_events) == 1
    assert final_response is not None
    assert final_response.finish_reason == "approval_needed"
    assert final_response.needs_approval is True
    assert len(final_response.approval_requests) == 1


@pytest.mark.asyncio
async def test_tool_approval_resume_after_approve() -> None:
    client = FakeModelClient([_tool_call_result(), _final_result("result is 5")])
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[sensitive_add],
    )

    first_response = await agent.run(UserMessage(content="calculate", source="user"))
    assert first_response.finish_reason == "approval_needed"
    assert first_response.context.waiting_for_approval is True

    first_response.context.approve_tool("call-1", reason="approved")

    resumed_response = await agent.run(context=first_response.context)
    assert resumed_response.finish_reason == "stop"
    assert any(
        message.role == "tool" and message.success
        for message in resumed_response.messages
    )
    assert resumed_response.messages[-1].content == "result is 5"


@pytest.mark.asyncio
async def test_tool_approval_resume_after_deny() -> None:
    client = FakeModelClient(
        [_tool_call_result(), _final_result("denial acknowledged")]
    )
    agent = Agent(
        name="assistant",
        description="desc",
        instructions="help",
        model_client=client,
        tools=[sensitive_add],
    )

    first_response = await agent.run(UserMessage(content="calculate", source="user"))
    assert first_response.finish_reason == "approval_needed"

    first_response.context.deny_tool("call-1", reason="unsafe operation")

    resumed_response = await agent.run(context=first_response.context)
    assert resumed_response.finish_reason == "stop"
    assert any(
        message.role == "tool"
        and not message.success
        and message.error == "Approval denied"
        for message in resumed_response.messages
    )
    assert resumed_response.messages[-1].content == "denial acknowledged"
