# Copyright (c) Microsoft. All rights reserved.

"""Unit tests for DurableAgentExecutor implementations.

Focuses on critical behavioral flows for executor strategies.
Run with: pytest tests/test_executors.py -v
"""

import time
from typing import Any
from unittest.mock import Mock

import pytest
from agent_framework import AgentRunResponse, Role
from durabletask.entities import EntityInstanceId
from durabletask.task import Task
from pydantic import BaseModel

from agent_framework_durabletask import DurableAgentThread
from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
from agent_framework_durabletask._executors import (
    ClientAgentExecutor,
    DurableAgentTask,
    OrchestrationAgentExecutor,
)
from agent_framework_durabletask._models import AgentSessionId, RunRequest


# Fixtures
@pytest.fixture
def mock_client() -> Mock:
    """Provide a mock client for ClientAgentExecutor tests."""
    client = Mock()
    client.signal_entity = Mock()
    client.get_entity = Mock(return_value=None)
    return client


@pytest.fixture
def mock_entity_task() -> Mock:
    """Provide a mock entity task."""
    return Mock(spec=Task)


@pytest.fixture
def mock_orchestration_context(mock_entity_task: Mock) -> Mock:
    """Provide a mock orchestration context with call_entity configured."""
    context = Mock()
    context.call_entity = Mock(return_value=mock_entity_task)
    return context


@pytest.fixture
def sample_run_request() -> RunRequest:
    """Provide a sample RunRequest for tests."""
    return RunRequest(message="test message", correlation_id="test-123")


@pytest.fixture
def client_executor(mock_client: Mock) -> ClientAgentExecutor:
    """Provide a ClientAgentExecutor with minimal polling for fast tests."""
    return ClientAgentExecutor(mock_client, max_poll_retries=1, poll_interval_seconds=0.01)


@pytest.fixture
def orchestration_executor(mock_orchestration_context: Mock) -> OrchestrationAgentExecutor:
    """Provide an OrchestrationAgentExecutor."""
    return OrchestrationAgentExecutor(mock_orchestration_context)


@pytest.fixture
def successful_agent_response() -> dict[str, Any]:
    """Provide a successful agent response dictionary."""
    return {
        "messages": [{"role": "assistant", "contents": [{"type": "text", "text": "Hello!"}]}],
        "created_at": "2025-12-30T10:00:00Z",
    }


class TestExecutorThreadCreation:
    """Test that executors properly create DurableAgentThread with parameters."""

    def test_client_executor_creates_durable_thread(self, mock_client: Mock) -> None:
        """Verify ClientAgentExecutor creates DurableAgentThread instances."""
        executor = ClientAgentExecutor(mock_client)

        thread = executor.get_new_thread("test_agent")

        assert isinstance(thread, DurableAgentThread)

    def test_client_executor_forwards_kwargs_to_thread(self, mock_client: Mock) -> None:
        """Verify ClientAgentExecutor forwards kwargs to DurableAgentThread creation."""
        executor = ClientAgentExecutor(mock_client)

        thread = executor.get_new_thread("test_agent", service_thread_id="client-123")

        assert isinstance(thread, DurableAgentThread)
        assert thread.service_thread_id == "client-123"

    def test_orchestration_executor_creates_durable_thread(
        self, orchestration_executor: OrchestrationAgentExecutor
    ) -> None:
        """Verify OrchestrationAgentExecutor creates DurableAgentThread instances."""
        thread = orchestration_executor.get_new_thread("test_agent")

        assert isinstance(thread, DurableAgentThread)

    def test_orchestration_executor_forwards_kwargs_to_thread(
        self, orchestration_executor: OrchestrationAgentExecutor
    ) -> None:
        """Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentThread creation."""
        thread = orchestration_executor.get_new_thread("test_agent", service_thread_id="orch-456")

        assert isinstance(thread, DurableAgentThread)
        assert thread.service_thread_id == "orch-456"


class TestClientAgentExecutorRun:
    """Test that ClientAgentExecutor.run_durable_agent works as implemented."""

    def test_client_executor_run_returns_response(
        self, client_executor: ClientAgentExecutor, sample_run_request: RunRequest
    ) -> None:
        """Verify ClientAgentExecutor.run_durable_agent returns AgentRunResponse (synchronous)."""
        result = client_executor.run_durable_agent("test_agent", sample_run_request)

        # Verify it returns an AgentRunResponse (synchronous, not a coroutine)
        assert isinstance(result, AgentRunResponse)
        assert result is not None


class TestClientAgentExecutorPollingConfiguration:
    """Test polling configuration parameters for ClientAgentExecutor."""

    def test_executor_uses_default_polling_parameters(self, mock_client: Mock) -> None:
        """Verify executor initializes with default polling parameters."""
        executor = ClientAgentExecutor(mock_client)

        assert executor.max_poll_retries == DEFAULT_MAX_POLL_RETRIES
        assert executor.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS

    def test_executor_accepts_custom_polling_parameters(self, mock_client: Mock) -> None:
        """Verify executor accepts and stores custom polling parameters."""
        executor = ClientAgentExecutor(mock_client, max_poll_retries=20, poll_interval_seconds=0.5)

        assert executor.max_poll_retries == 20
        assert executor.poll_interval_seconds == 0.5

    def test_executor_respects_custom_max_poll_retries(self, mock_client: Mock, sample_run_request: RunRequest) -> None:
        """Verify executor respects custom max_poll_retries during polling."""
        # Create executor with only 2 retries
        executor = ClientAgentExecutor(mock_client, max_poll_retries=2, poll_interval_seconds=0.01)

        # Run the agent
        result = executor.run_durable_agent("test_agent", sample_run_request)

        # Verify it returns AgentRunResponse (should timeout after 2 attempts)
        assert isinstance(result, AgentRunResponse)

        # Verify get_entity was called 2 times (max_poll_retries)
        assert mock_client.get_entity.call_count == 2

    def test_executor_respects_custom_poll_interval(self, mock_client: Mock, sample_run_request: RunRequest) -> None:
        """Verify executor respects custom poll_interval_seconds during polling."""
        # Create executor with very short interval
        executor = ClientAgentExecutor(mock_client, max_poll_retries=3, poll_interval_seconds=0.01)

        # Measure time taken
        start = time.time()
        result = executor.run_durable_agent("test_agent", sample_run_request)
        elapsed = time.time() - start

        # Should take roughly 3 * 0.01 = 0.03 seconds (plus overhead)
        # Be generous with timing to avoid flakiness
        assert elapsed < 0.2  # Should be quick with 0.01 interval
        assert isinstance(result, AgentRunResponse)


class TestOrchestrationAgentExecutorRun:
    """Test OrchestrationAgentExecutor.run_durable_agent implementation."""

    def test_orchestration_executor_run_returns_durable_agent_task(
        self, orchestration_executor: OrchestrationAgentExecutor, sample_run_request: RunRequest
    ) -> None:
        """Verify OrchestrationAgentExecutor.run_durable_agent returns DurableAgentTask."""
        result = orchestration_executor.run_durable_agent("test_agent", sample_run_request)

        assert isinstance(result, DurableAgentTask)

    def test_orchestration_executor_calls_entity_with_correct_parameters(
        self,
        mock_orchestration_context: Mock,
        orchestration_executor: OrchestrationAgentExecutor,
        sample_run_request: RunRequest,
    ) -> None:
        """Verify call_entity is invoked with correct entity ID and request."""
        orchestration_executor.run_durable_agent("test_agent", sample_run_request)

        # Verify call_entity was called once
        assert mock_orchestration_context.call_entity.call_count == 1

        # Get the call arguments
        call_args = mock_orchestration_context.call_entity.call_args
        entity_id_arg = call_args[0][0]
        operation_arg = call_args[0][1]
        request_dict_arg = call_args[0][2]

        # Verify entity ID
        assert isinstance(entity_id_arg, EntityInstanceId)
        assert entity_id_arg.entity == "dafx-test_agent"

        # Verify operation name
        assert operation_arg == "run"

        # Verify request dict
        assert request_dict_arg == sample_run_request.to_dict()

    def test_orchestration_executor_uses_thread_session_id(
        self,
        mock_orchestration_context: Mock,
        orchestration_executor: OrchestrationAgentExecutor,
        sample_run_request: RunRequest,
    ) -> None:
        """Verify executor uses thread's session ID when provided."""
        # Create thread with specific session ID
        session_id = AgentSessionId(name="test_agent", key="specific-key-123")
        thread = DurableAgentThread.from_session_id(session_id)

        result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, thread=thread)

        # Verify call_entity was called with the specific key
        call_args = mock_orchestration_context.call_entity.call_args
        entity_id_arg = call_args[0][0]

        assert entity_id_arg.key == "specific-key-123"
        assert isinstance(result, DurableAgentTask)


class TestDurableAgentTask:
    """Test DurableAgentTask completion and response transformation."""

    def test_durable_agent_task_transforms_successful_result(
        self, mock_entity_task: Mock, successful_agent_response: dict[str, Any]
    ) -> None:
        """Verify DurableAgentTask converts successful entity result to AgentRunResponse."""
        mock_entity_task.is_failed = False
        mock_entity_task.get_result = Mock(return_value=successful_agent_response)

        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")

        # Simulate child task completion
        task.on_child_completed(mock_entity_task)

        assert task.is_complete
        result = task.get_result()
        assert isinstance(result, AgentRunResponse)
        assert len(result.messages) == 1
        assert result.messages[0].role == Role.ASSISTANT

    def test_durable_agent_task_propagates_failure(self, mock_entity_task: Mock) -> None:
        """Verify DurableAgentTask propagates task failures."""
        mock_entity_task.is_failed = True
        mock_entity_task.get_exception = Mock(return_value=ValueError("Entity error"))

        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")

        # Simulate child task completion with failure
        task.on_child_completed(mock_entity_task)

        assert task.is_complete
        assert task.is_failed
        exception = task.get_exception()
        assert isinstance(exception, ValueError)
        assert str(exception) == "Entity error"

    def test_durable_agent_task_validates_response_format(self, mock_entity_task: Mock) -> None:
        """Verify DurableAgentTask validates response format when provided."""
        mock_entity_task.is_failed = False
        mock_entity_task.get_result = Mock(
            return_value={
                "messages": [{"role": "assistant", "contents": [{"type": "text", "text": '{"answer": "42"}'}]}],
                "created_at": "2025-12-30T10:00:00Z",
            }
        )

        class TestResponse(BaseModel):
            answer: str

        task = DurableAgentTask(entity_task=mock_entity_task, response_format=TestResponse, correlation_id="test-123")

        # Simulate child task completion
        task.on_child_completed(mock_entity_task)

        assert task.is_complete
        result = task.get_result()
        assert isinstance(result, AgentRunResponse)

    def test_durable_agent_task_ignores_duplicate_completion(
        self, mock_entity_task: Mock, successful_agent_response: dict[str, Any]
    ) -> None:
        """Verify DurableAgentTask ignores duplicate completion calls."""
        mock_entity_task.is_failed = False
        mock_entity_task.get_result = Mock(return_value=successful_agent_response)

        task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")

        # Simulate child task completion twice
        task.on_child_completed(mock_entity_task)
        first_result = task.get_result()

        task.on_child_completed(mock_entity_task)
        second_result = task.get_result()

        # Should be the same result, get_result should only be called once
        assert first_result is second_result
        assert mock_entity_task.get_result.call_count == 1


if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])