# Copyright (c) Microsoft. All rights reserved.

"""Unit tests for orchestration request info support."""

from collections.abc import AsyncIterable
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock

import pytest
from agent_framework import (
    AgentResponse,
    AgentResponseUpdate,
    AgentSession,
    Content,
    Message,
    SupportsAgentRun,
)
from agent_framework._workflows._agent_executor import AgentExecutorRequest, AgentExecutorResponse
from agent_framework._workflows._workflow_context import WorkflowContext

from agent_framework_orchestrations._orchestration_request_info import (
    AgentApprovalExecutor,
    AgentRequestInfoExecutor,
    AgentRequestInfoResponse,
    resolve_request_info_filter,
)


class TestResolveRequestInfoFilter:
    """Tests for resolve_request_info_filter function."""

    def test_returns_empty_set_for_none_input(self):
        """Test that None input returns empty set (no filtering)."""
        result = resolve_request_info_filter(None)
        assert result == set()

    def test_returns_empty_set_for_empty_list(self):
        """Test that empty list returns empty set."""
        result = resolve_request_info_filter([])
        assert result == set()

    def test_resolves_string_names(self):
        """Test resolving string agent names."""
        result = resolve_request_info_filter(["agent1", "agent2"])
        assert result == {"agent1", "agent2"}

    def test_resolves_agent_display_names(self):
        """Test resolving SupportsAgentRun instances by name attribute."""
        agent1 = MagicMock(spec=SupportsAgentRun)
        agent1.name = "writer"
        agent2 = MagicMock(spec=SupportsAgentRun)
        agent2.name = "reviewer"

        result = resolve_request_info_filter([agent1, agent2])
        assert result == {"writer", "reviewer"}

    def test_mixed_types(self):
        """Test resolving a mix of strings and agents."""
        agent = MagicMock(spec=SupportsAgentRun)
        agent.name = "writer"

        result = resolve_request_info_filter(["manual_name", agent])
        assert result == {"manual_name", "writer"}

    def test_raises_on_unsupported_type(self):
        """Test that unsupported types raise TypeError."""
        with pytest.raises(TypeError, match="Unsupported type for request_info filter"):
            resolve_request_info_filter([123])  # type: ignore


class TestAgentRequestInfoResponse:
    """Tests for AgentRequestInfoResponse dataclass."""

    def test_create_response_with_messages(self):
        """Test creating an AgentRequestInfoResponse with messages."""
        messages = [Message(role="user", contents=["Additional info"])]
        response = AgentRequestInfoResponse(messages=messages)

        assert response.messages == messages

    def test_from_messages_factory(self):
        """Test creating response from Message list."""
        messages = [
            Message(role="user", contents=["Message 1"]),
            Message(role="user", contents=["Message 2"]),
        ]
        response = AgentRequestInfoResponse.from_messages(messages)

        assert response.messages == messages

    def test_from_strings_factory(self):
        """Test creating response from string list."""
        texts = ["First message", "Second message"]
        response = AgentRequestInfoResponse.from_strings(texts)

        assert len(response.messages) == 2
        assert response.messages[0].role == "user"
        assert response.messages[0].text == "First message"
        assert response.messages[1].role == "user"
        assert response.messages[1].text == "Second message"

    def test_approve_factory(self):
        """Test creating an approval response (empty messages)."""
        response = AgentRequestInfoResponse.approve()

        assert response.messages == []


class TestAgentRequestInfoExecutor:
    """Tests for AgentRequestInfoExecutor."""

    @pytest.mark.asyncio
    async def test_request_info_handler(self):
        """Test that request_info handler calls ctx.request_info."""
        executor = AgentRequestInfoExecutor(id="test_executor")

        agent_response = AgentResponse(messages=[Message(role="assistant", contents=["Agent response"])])
        executor_response = AgentExecutorResponse(
            executor_id="test_agent",
            agent_response=agent_response,
            full_conversation=agent_response.messages,
        )

        ctx = MagicMock(spec=WorkflowContext)
        ctx.request_info = AsyncMock()

        await executor.request_info(executor_response, ctx)

        ctx.request_info.assert_called_once_with(executor_response, AgentRequestInfoResponse)

    @pytest.mark.asyncio
    async def test_handle_request_info_response_with_messages(self):
        """Test response handler when user provides additional messages."""
        executor = AgentRequestInfoExecutor(id="test_executor")

        agent_response = AgentResponse(messages=[Message(role="assistant", contents=["Original"])])
        original_request = AgentExecutorResponse(
            executor_id="test_agent",
            agent_response=agent_response,
            full_conversation=agent_response.messages,
        )

        response = AgentRequestInfoResponse.from_strings(["Additional input"])

        ctx = MagicMock(spec=WorkflowContext)
        ctx.send_message = AsyncMock()

        await executor.handle_request_info_response(original_request, response, ctx)

        # Should send new request with additional messages
        ctx.send_message.assert_called_once()
        call_args = ctx.send_message.call_args[0][0]
        assert isinstance(call_args, AgentExecutorRequest)
        assert call_args.should_respond is True
        assert len(call_args.messages) == 1
        assert call_args.messages[0].text == "Additional input"

    @pytest.mark.asyncio
    async def test_handle_request_info_response_approval(self):
        """Test response handler when user approves (no additional messages)."""
        executor = AgentRequestInfoExecutor(id="test_executor")

        agent_response = AgentResponse(messages=[Message(role="assistant", contents=["Original"])])
        original_request = AgentExecutorResponse(
            executor_id="test_agent",
            agent_response=agent_response,
            full_conversation=agent_response.messages,
        )

        response = AgentRequestInfoResponse.approve()

        ctx = MagicMock(spec=WorkflowContext)
        ctx.yield_output = AsyncMock()

        await executor.handle_request_info_response(original_request, response, ctx)

        # Should yield original response without modification
        ctx.yield_output.assert_called_once_with(original_request)


class _TestAgent:
    """Simple test agent implementation."""

    def __init__(self, id: str, name: str | None = None, description: str | None = None):
        self._id = id
        self._name = name
        self._description = description

    @property
    def id(self) -> str:
        return self._id

    @property
    def name(self) -> str | None:
        return self._name

    @property
    def display_name(self) -> str:
        return self._name or self._id

    @property
    def description(self) -> str | None:
        return self._description

    async def run(
        self,
        messages: str | Message | list[str] | list[Message] | None = None,
        *,
        stream: bool = False,
        session: AgentSession | None = None,
        **kwargs: Any,
    ) -> Any:
        """Dummy run method."""
        if stream:
            return self._run_stream_impl()
        return AgentResponse(messages=[Message(role="assistant", contents=["Test response"])])

    async def _run_stream_impl(self) -> AsyncIterable[AgentResponseUpdate]:
        yield AgentResponseUpdate(contents=[Content.from_text(text="Test response stream")])

    def create_session(self, **kwargs: Any) -> AgentSession:
        """Creates a new conversation session for the agent."""
        return AgentSession(**kwargs)

    def get_session(self, service_session_id: str, *, session_id: str | None = None) -> AgentSession:
        """Gets a conversation session for the agent."""
        return AgentSession(service_session_id=service_session_id, session_id=session_id)


class TestAgentApprovalExecutor:
    """Tests for AgentApprovalExecutor."""

    def test_initialization(self):
        """Test that AgentApprovalExecutor initializes correctly."""
        agent = _TestAgent(id="test_id", name="test_agent", description="Test agent description")

        executor = AgentApprovalExecutor(cast(SupportsAgentRun, agent))

        assert executor.id == "test_agent"
        assert executor.description == "Test agent description"

    def test_builds_workflow_with_agent_and_request_info_executors(self):
        """Test that the internal workflow is created successfully."""
        agent = _TestAgent(id="test_id", name="test_agent", description="Test description")

        executor = AgentApprovalExecutor(cast(SupportsAgentRun, agent))

        # Verify the executor has a workflow
        assert executor.workflow is not None
        assert executor.id == "test_agent"

    def test_propagate_request_enabled(self):
        """Test that AgentApprovalExecutor has propagate_request enabled."""
        agent = _TestAgent(id="test_id", name="test_agent", description="Test description")

        executor = AgentApprovalExecutor(cast(SupportsAgentRun, agent))

        assert executor._propagate_request is True  # type: ignore
