"""Tests for MCP task-augmented execution support in MCPClient."""

import asyncio
from datetime import timedelta
from unittest.mock import AsyncMock, MagicMock

import pytest
from mcp import ListToolsResult
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import TextContent as MCPTextContent
from mcp.types import Tool as MCPTool
from mcp.types import ToolExecution

from strands.tools.mcp import MCPClient, TasksConfig
from strands.tools.mcp.mcp_tasks import DEFAULT_TASK_POLL_TIMEOUT, DEFAULT_TASK_TTL

from .conftest import create_server_capabilities


class TestTasksOptIn:
    """Tests for task opt-in behavior via tasks config."""

    @pytest.mark.parametrize(
        "tasks_config,expected_enabled",
        [
            (None, False),
            ({}, True),
        ],
    )
    def test_tasks_enabled_state(self, mock_transport, mock_session, tasks_config, expected_enabled):
        """Test _is_tasks_enabled based on tasks config."""
        with MCPClient(mock_transport["transport_callable"], tasks_config=tasks_config) as client:
            assert client._is_tasks_enabled() is expected_enabled

    def test_should_use_task_requires_opt_in(self, mock_transport, mock_session):
        """Test that _should_use_task returns False without opt-in even with server/tool support."""
        with MCPClient(mock_transport["transport_callable"]) as client:
            client._server_task_capable = True
            assert client._should_use_task("test_tool") is False

        with MCPClient(mock_transport["transport_callable"], tasks_config={}) as client:
            client._server_task_capable = True
            client._tool_task_support_cache["test_tool"] = "required"
            assert client._should_use_task("test_tool") is True


class TestTaskConfiguration:
    """Tests for task-related configuration options."""

    @pytest.mark.parametrize(
        "config,expected_ttl,expected_timeout",
        [
            ({}, DEFAULT_TASK_TTL, DEFAULT_TASK_POLL_TIMEOUT),
            ({"ttl": timedelta(seconds=120)}, timedelta(seconds=120), DEFAULT_TASK_POLL_TIMEOUT),
            ({"poll_timeout": timedelta(seconds=60)}, DEFAULT_TASK_TTL, timedelta(seconds=60)),
            (
                {"ttl": timedelta(seconds=120), "poll_timeout": timedelta(seconds=60)},
                timedelta(seconds=120),
                timedelta(seconds=60),
            ),
        ],
    )
    def test_task_config_values(self, mock_transport, mock_session, config, expected_ttl, expected_timeout):
        """Test task configuration values with various configs."""
        with MCPClient(mock_transport["transport_callable"], tasks_config=config) as client:
            config_actual = client._get_task_config()
            assert config_actual.get("ttl") == expected_ttl
            assert config_actual.get("poll_timeout") == expected_timeout

    def test_stop_resets_task_caches(self, mock_transport, mock_session):
        """Test that stop() resets the task support caches."""
        with MCPClient(mock_transport["transport_callable"], tasks_config={}) as client:
            client._server_task_capable = True
            client._tool_task_support_cache["tool1"] = "required"
        assert client._server_task_capable is None
        assert client._tool_task_support_cache == {}


class TestTaskExecution:
    """Tests for task execution and error handling."""

    def _setup_task_tool(self, mock_session, tool_name: str) -> None:
        """Helper to set up a mock task-enabled tool."""
        mock_session.get_server_capabilities = MagicMock(return_value=create_server_capabilities(True))
        mock_tool = MCPTool(
            name=tool_name,
            description="A test tool",
            inputSchema={"type": "object"},
            execution=ToolExecution(taskSupport="optional"),
        )
        mock_session.list_tools = AsyncMock(return_value=ListToolsResult(tools=[mock_tool], nextCursor=None))
        mock_create_result = MagicMock()
        mock_create_result.task.taskId = "test-task-id"
        mock_session.experimental = MagicMock()
        mock_session.experimental.call_tool_as_task = AsyncMock(return_value=mock_create_result)

    @pytest.mark.parametrize(
        "status,status_message,expected_text",
        [
            ("failed", "Something went wrong", "Something went wrong"),
            ("cancelled", None, "cancelled"),
            ("unknown_status", None, "unexpected task status"),
        ],
    )
    def test_terminal_status_handling(self, mock_transport, mock_session, status, status_message, expected_text):
        """Test handling of terminal task statuses."""
        mock_create_result = MagicMock()
        mock_create_result.task.taskId = f"task-{status}"
        mock_session.experimental.call_tool_as_task = AsyncMock(return_value=mock_create_result)

        async def mock_poll_task(task_id):
            yield MagicMock(status=status, statusMessage=status_message)

        mock_session.experimental.poll_task = mock_poll_task

        with MCPClient(mock_transport["transport_callable"], tasks_config=TasksConfig()) as client:
            client._server_task_capable = True
            client._tool_task_support_cache["test_tool"] = "required"
            result = client.call_tool_sync(tool_use_id="test-id", name="test_tool", arguments={})
            assert result["status"] == "error"
            assert expected_text.lower() in result["content"][0].get("text", "").lower()

    @pytest.mark.asyncio
    async def test_polling_timeout(self, mock_transport, mock_session):
        """Test that task polling times out properly."""
        self._setup_task_tool(mock_session, "slow_tool")

        async def infinite_poll(task_id):
            while True:
                await asyncio.sleep(1)
                yield MagicMock(status="running")

        mock_session.experimental.poll_task = infinite_poll

        with MCPClient(
            mock_transport["transport_callable"], tasks_config=TasksConfig(poll_timeout=timedelta(seconds=0.1))
        ) as client:
            client.list_tools_sync()
            result = await client.call_tool_async(tool_use_id="t", name="slow_tool", arguments={})
            assert result["status"] == "error"
            assert "timed out" in result["content"][0].get("text", "").lower()

    @pytest.mark.asyncio
    async def test_explicit_timeout_overrides_default(self, mock_transport, mock_session):
        """Test that read_timeout_seconds overrides the default poll timeout."""
        self._setup_task_tool(mock_session, "timeout_tool")

        async def infinite_poll(task_id):
            while True:
                await asyncio.sleep(1)
                yield MagicMock(status="running")

        mock_session.experimental.poll_task = infinite_poll

        with MCPClient(
            mock_transport["transport_callable"], tasks_config=TasksConfig(poll_timeout=timedelta(minutes=5))
        ) as client:
            client.list_tools_sync()
            result = await client.call_tool_async(
                tool_use_id="t", name="timeout_tool", arguments={}, read_timeout_seconds=timedelta(seconds=0.1)
            )
            assert result["status"] == "error"
            assert "timed out" in result["content"][0].get("text", "").lower()

    @pytest.mark.asyncio
    async def test_result_retrieval_failure(self, mock_transport, mock_session):
        """Test that get_task_result failures are handled gracefully."""
        self._setup_task_tool(mock_session, "failing_tool")

        async def successful_poll(task_id):
            yield MagicMock(status="completed", statusMessage=None)

        mock_session.experimental.poll_task = successful_poll
        mock_session.experimental.get_task_result = AsyncMock(side_effect=Exception("Network error"))

        with MCPClient(mock_transport["transport_callable"], tasks_config=TasksConfig()) as client:
            client.list_tools_sync()
            result = await client.call_tool_async(tool_use_id="t", name="failing_tool", arguments={})
            assert result["status"] == "error"
            assert "result retrieval failed" in result["content"][0].get("text", "").lower()

    @pytest.mark.asyncio
    async def test_empty_poll_result(self, mock_transport, mock_session):
        """Test handling when poll_task yields nothing."""
        self._setup_task_tool(mock_session, "empty_poll_tool")

        async def empty_poll(task_id):
            return
            yield  # noqa: B901

        mock_session.experimental.poll_task = empty_poll

        with MCPClient(mock_transport["transport_callable"], tasks_config=TasksConfig()) as client:
            client.list_tools_sync()
            result = await client.call_tool_async(tool_use_id="t", name="empty_poll_tool", arguments={})
            assert result["status"] == "error"
            assert "without status" in result["content"][0].get("text", "").lower()

    @pytest.mark.asyncio
    async def test_successful_completion(self, mock_transport, mock_session):
        """Test successful task completion."""
        self._setup_task_tool(mock_session, "success_tool")

        async def poll(task_id):
            yield MagicMock(status="completed", statusMessage=None)

        mock_session.experimental.poll_task = poll
        mock_session.experimental.get_task_result = AsyncMock(
            return_value=MCPCallToolResult(content=[MCPTextContent(type="text", text="Done")], isError=False)
        )

        with MCPClient(mock_transport["transport_callable"], tasks_config=TasksConfig()) as client:
            client.list_tools_sync()
            result = await client.call_tool_async(tool_use_id="t", name="success_tool", arguments={})
            assert result["status"] == "success"
            assert "Done" in result["content"][0].get("text", "")


class TestTaskMetaForwarding:
    """Tests for meta parameter forwarding in task-augmented execution."""

    def _setup_task_tool_with_meta(self, mock_session, tool_name: str) -> MagicMock:
        """Helper to set up a mock task-enabled tool and return the experimental mock."""
        mock_session.get_server_capabilities = MagicMock(return_value=create_server_capabilities(True))
        mock_tool = MCPTool(
            name=tool_name,
            description="A test tool",
            inputSchema={"type": "object"},
            execution=ToolExecution(taskSupport="optional"),
        )
        mock_session.list_tools = AsyncMock(return_value=ListToolsResult(tools=[mock_tool], nextCursor=None))
        mock_create_result = MagicMock()
        mock_create_result.task.taskId = "test-task-id"
        mock_session.experimental = MagicMock()
        mock_session.experimental.call_tool_as_task = AsyncMock(return_value=mock_create_result)

        async def successful_poll(task_id):
            yield MagicMock(status="completed", statusMessage=None)

        mock_session.experimental.poll_task = successful_poll
        mock_session.experimental.get_task_result = AsyncMock(
            return_value=MCPCallToolResult(content=[MCPTextContent(type="text", text="Done")], isError=False)
        )

        return mock_session.experimental

    def test_call_tool_sync_forwards_meta_to_task(self, mock_transport, mock_session):
        """Test that call_tool_sync forwards meta to call_tool_as_task."""
        experimental = self._setup_task_tool_with_meta(mock_session, "meta_tool")
        meta = {"com.example/request_id": "abc-123"}

        with MCPClient(mock_transport["transport_callable"], tasks_config=TasksConfig()) as client:
            client.list_tools_sync()
            client.call_tool_sync(tool_use_id="test-id", name="meta_tool", arguments={"param": "value"}, meta=meta)

            experimental.call_tool_as_task.assert_called_once()
            call_kwargs = experimental.call_tool_as_task.call_args
            assert call_kwargs.kwargs.get("meta") == meta

    @pytest.mark.asyncio
    async def test_call_tool_async_forwards_meta_to_task(self, mock_transport, mock_session):
        """Test that call_tool_async forwards meta to call_tool_as_task."""
        experimental = self._setup_task_tool_with_meta(mock_session, "meta_tool")
        meta = {"com.example/trace_id": "xyz-456"}

        with MCPClient(mock_transport["transport_callable"], tasks_config=TasksConfig()) as client:
            client.list_tools_sync()
            await client.call_tool_async(
                tool_use_id="test-id", name="meta_tool", arguments={"param": "value"}, meta=meta
            )

            experimental.call_tool_as_task.assert_called_once()
            call_kwargs = experimental.call_tool_as_task.call_args
            assert call_kwargs.kwargs.get("meta") == meta

    def test_call_tool_sync_forwards_none_meta_to_task(self, mock_transport, mock_session):
        """Test that call_tool_sync forwards None meta to call_tool_as_task when not provided."""
        experimental = self._setup_task_tool_with_meta(mock_session, "no_meta_tool")

        with MCPClient(mock_transport["transport_callable"], tasks_config=TasksConfig()) as client:
            client.list_tools_sync()
            client.call_tool_sync(tool_use_id="test-id", name="no_meta_tool", arguments={"param": "value"})

            experimental.call_tool_as_task.assert_called_once()
            call_kwargs = experimental.call_tool_as_task.call_args
            assert call_kwargs.kwargs.get("meta") is None
