"""Tests for the experimental client task methods (session.experimental)."""

from dataclasses import dataclass, field
from typing import Any

import anyio
import pytest
from anyio import Event
from anyio.abc import TaskGroup

from mcp.client.session import ClientSession
from mcp.server import Server
from mcp.server.lowlevel import NotificationOptions
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.shared.experimental.tasks.helpers import task_execution
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from mcp.types import (
    CallToolRequest,
    CallToolRequestParams,
    CallToolResult,
    CancelTaskRequest,
    CancelTaskResult,
    ClientRequest,
    ClientResult,
    CreateTaskResult,
    GetTaskPayloadRequest,
    GetTaskPayloadResult,
    GetTaskRequest,
    GetTaskResult,
    ListTasksRequest,
    ListTasksResult,
    ServerNotification,
    ServerRequest,
    TaskMetadata,
    TextContent,
    Tool,
)


@dataclass
class AppContext:
    """Application context passed via lifespan_context."""

    task_group: TaskGroup
    store: InMemoryTaskStore
    task_done_events: dict[str, Event] = field(default_factory=lambda: {})


@pytest.mark.anyio
async def test_session_experimental_get_task() -> None:
    """Test session.experimental.get_task() method."""
    # Note: We bypass the normal lifespan mechanism
    server: Server[AppContext, Any] = Server("test-server")  # type: ignore[assignment]
    store = InMemoryTaskStore()

    @server.list_tools()
    async def list_tools():
        return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
        ctx = server.request_context
        app = ctx.lifespan_context
        if ctx.experimental.is_task:
            task_metadata = ctx.experimental.task_metadata
            assert task_metadata is not None
            task = await app.store.create_task(task_metadata)

            done_event = Event()
            app.task_done_events[task.taskId] = done_event

            async def do_work():
                async with task_execution(task.taskId, app.store) as task_ctx:
                    await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")]))
                done_event.set()

            app.task_group.start_soon(do_work)
            return CreateTaskResult(task=task)

        raise NotImplementedError

    @server.experimental.get_task()
    async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
        app = server.request_context.lifespan_context
        task = await app.store.get_task(request.params.taskId)
        assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
        return GetTaskResult(
            taskId=task.taskId,
            status=task.status,
            statusMessage=task.statusMessage,
            createdAt=task.createdAt,
            lastUpdatedAt=task.lastUpdatedAt,
            ttl=task.ttl,
            pollInterval=task.pollInterval,
        )

    # Set up streams
    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def message_handler(
        message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
    ) -> None: ...  # pragma: no branch

    async def run_server(app_context: AppContext):
        async with ServerSession(
            client_to_server_receive,
            server_to_client_send,
            InitializationOptions(
                server_name="test-server",
                server_version="1.0.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        ) as server_session:
            async for message in server_session.incoming_messages:
                await server._handle_message(message, server_session, app_context, raise_exceptions=False)

    async with anyio.create_task_group() as tg:
        app_context = AppContext(task_group=tg, store=store)
        tg.start_soon(run_server, app_context)

        async with ClientSession(
            server_to_client_receive,
            client_to_server_send,
            message_handler=message_handler,
        ) as client_session:
            await client_session.initialize()

            # Create a task
            create_result = await client_session.send_request(
                ClientRequest(
                    CallToolRequest(
                        params=CallToolRequestParams(
                            name="test_tool",
                            arguments={},
                            task=TaskMetadata(ttl=60000),
                        )
                    )
                ),
                CreateTaskResult,
            )
            task_id = create_result.task.taskId

            # Wait for task to complete
            await app_context.task_done_events[task_id].wait()

            # Use session.experimental to get task status
            task_status = await client_session.experimental.get_task(task_id)

            assert task_status.taskId == task_id
            assert task_status.status == "completed"

            tg.cancel_scope.cancel()


@pytest.mark.anyio
async def test_session_experimental_get_task_result() -> None:
    """Test session.experimental.get_task_result() method."""
    server: Server[AppContext, Any] = Server("test-server")  # type: ignore[assignment]
    store = InMemoryTaskStore()

    @server.list_tools()
    async def list_tools():
        return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
        ctx = server.request_context
        app = ctx.lifespan_context
        if ctx.experimental.is_task:
            task_metadata = ctx.experimental.task_metadata
            assert task_metadata is not None
            task = await app.store.create_task(task_metadata)

            done_event = Event()
            app.task_done_events[task.taskId] = done_event

            async def do_work():
                async with task_execution(task.taskId, app.store) as task_ctx:
                    await task_ctx.complete(
                        CallToolResult(content=[TextContent(type="text", text="Task result content")])
                    )
                done_event.set()

            app.task_group.start_soon(do_work)
            return CreateTaskResult(task=task)

        raise NotImplementedError

    @server.experimental.get_task_result()
    async def handle_get_task_result(
        request: GetTaskPayloadRequest,
    ) -> GetTaskPayloadResult:
        app = server.request_context.lifespan_context
        result = await app.store.get_result(request.params.taskId)
        assert result is not None, f"Test setup error: result for {request.params.taskId} should exist"
        assert isinstance(result, CallToolResult)
        return GetTaskPayloadResult(**result.model_dump())

    # Set up streams
    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def message_handler(
        message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
    ) -> None: ...  # pragma: no branch

    async def run_server(app_context: AppContext):
        async with ServerSession(
            client_to_server_receive,
            server_to_client_send,
            InitializationOptions(
                server_name="test-server",
                server_version="1.0.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        ) as server_session:
            async for message in server_session.incoming_messages:
                await server._handle_message(message, server_session, app_context, raise_exceptions=False)

    async with anyio.create_task_group() as tg:
        app_context = AppContext(task_group=tg, store=store)
        tg.start_soon(run_server, app_context)

        async with ClientSession(
            server_to_client_receive,
            client_to_server_send,
            message_handler=message_handler,
        ) as client_session:
            await client_session.initialize()

            # Create a task
            create_result = await client_session.send_request(
                ClientRequest(
                    CallToolRequest(
                        params=CallToolRequestParams(
                            name="test_tool",
                            arguments={},
                            task=TaskMetadata(ttl=60000),
                        )
                    )
                ),
                CreateTaskResult,
            )
            task_id = create_result.task.taskId

            # Wait for task to complete
            await app_context.task_done_events[task_id].wait()

            # Use TaskClient to get task result
            task_result = await client_session.experimental.get_task_result(task_id, CallToolResult)

            assert len(task_result.content) == 1
            content = task_result.content[0]
            assert isinstance(content, TextContent)
            assert content.text == "Task result content"

            tg.cancel_scope.cancel()


@pytest.mark.anyio
async def test_session_experimental_list_tasks() -> None:
    """Test TaskClient.list_tasks() method."""
    server: Server[AppContext, Any] = Server("test-server")  # type: ignore[assignment]
    store = InMemoryTaskStore()

    @server.list_tools()
    async def list_tools():
        return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
        ctx = server.request_context
        app = ctx.lifespan_context
        if ctx.experimental.is_task:
            task_metadata = ctx.experimental.task_metadata
            assert task_metadata is not None
            task = await app.store.create_task(task_metadata)

            done_event = Event()
            app.task_done_events[task.taskId] = done_event

            async def do_work():
                async with task_execution(task.taskId, app.store) as task_ctx:
                    await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text="Done")]))
                done_event.set()

            app.task_group.start_soon(do_work)
            return CreateTaskResult(task=task)

        raise NotImplementedError

    @server.experimental.list_tasks()
    async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
        app = server.request_context.lifespan_context
        tasks_list, next_cursor = await app.store.list_tasks(cursor=request.params.cursor if request.params else None)
        return ListTasksResult(tasks=tasks_list, nextCursor=next_cursor)

    # Set up streams
    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def message_handler(
        message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
    ) -> None: ...  # pragma: no branch

    async def run_server(app_context: AppContext):
        async with ServerSession(
            client_to_server_receive,
            server_to_client_send,
            InitializationOptions(
                server_name="test-server",
                server_version="1.0.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        ) as server_session:
            async for message in server_session.incoming_messages:
                await server._handle_message(message, server_session, app_context, raise_exceptions=False)

    async with anyio.create_task_group() as tg:
        app_context = AppContext(task_group=tg, store=store)
        tg.start_soon(run_server, app_context)

        async with ClientSession(
            server_to_client_receive,
            client_to_server_send,
            message_handler=message_handler,
        ) as client_session:
            await client_session.initialize()

            # Create two tasks
            for _ in range(2):
                create_result = await client_session.send_request(
                    ClientRequest(
                        CallToolRequest(
                            params=CallToolRequestParams(
                                name="test_tool",
                                arguments={},
                                task=TaskMetadata(ttl=60000),
                            )
                        )
                    ),
                    CreateTaskResult,
                )
                await app_context.task_done_events[create_result.task.taskId].wait()

            # Use TaskClient to list tasks
            list_result = await client_session.experimental.list_tasks()

            assert len(list_result.tasks) == 2

            tg.cancel_scope.cancel()


@pytest.mark.anyio
async def test_session_experimental_cancel_task() -> None:
    """Test TaskClient.cancel_task() method."""
    server: Server[AppContext, Any] = Server("test-server")  # type: ignore[assignment]
    store = InMemoryTaskStore()

    @server.list_tools()
    async def list_tools():
        return [Tool(name="test_tool", description="Test", inputSchema={"type": "object"})]

    @server.call_tool()
    async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
        ctx = server.request_context
        app = ctx.lifespan_context
        if ctx.experimental.is_task:
            task_metadata = ctx.experimental.task_metadata
            assert task_metadata is not None
            task = await app.store.create_task(task_metadata)
            # Don't start any work - task stays in "working" status
            return CreateTaskResult(task=task)

        raise NotImplementedError

    @server.experimental.get_task()
    async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
        app = server.request_context.lifespan_context
        task = await app.store.get_task(request.params.taskId)
        assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
        return GetTaskResult(
            taskId=task.taskId,
            status=task.status,
            statusMessage=task.statusMessage,
            createdAt=task.createdAt,
            lastUpdatedAt=task.lastUpdatedAt,
            ttl=task.ttl,
            pollInterval=task.pollInterval,
        )

    @server.experimental.cancel_task()
    async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult:
        app = server.request_context.lifespan_context
        task = await app.store.get_task(request.params.taskId)
        assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
        await app.store.update_task(request.params.taskId, status="cancelled")
        # CancelTaskResult extends Task, so we need to return the updated task info
        updated_task = await app.store.get_task(request.params.taskId)
        assert updated_task is not None
        return CancelTaskResult(
            taskId=updated_task.taskId,
            status=updated_task.status,
            createdAt=updated_task.createdAt,
            lastUpdatedAt=updated_task.lastUpdatedAt,
            ttl=updated_task.ttl,
        )

    # Set up streams
    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)

    async def message_handler(
        message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
    ) -> None: ...  # pragma: no branch

    async def run_server(app_context: AppContext):
        async with ServerSession(
            client_to_server_receive,
            server_to_client_send,
            InitializationOptions(
                server_name="test-server",
                server_version="1.0.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        ) as server_session:
            async for message in server_session.incoming_messages:
                await server._handle_message(message, server_session, app_context, raise_exceptions=False)

    async with anyio.create_task_group() as tg:
        app_context = AppContext(task_group=tg, store=store)
        tg.start_soon(run_server, app_context)

        async with ClientSession(
            server_to_client_receive,
            client_to_server_send,
            message_handler=message_handler,
        ) as client_session:
            await client_session.initialize()

            # Create a task (but don't complete it)
            create_result = await client_session.send_request(
                ClientRequest(
                    CallToolRequest(
                        params=CallToolRequestParams(
                            name="test_tool",
                            arguments={},
                            task=TaskMetadata(ttl=60000),
                        )
                    )
                ),
                CreateTaskResult,
            )
            task_id = create_result.task.taskId

            # Verify task is working
            status_before = await client_session.experimental.get_task(task_id)
            assert status_before.status == "working"

            # Cancel the task
            await client_session.experimental.cancel_task(task_id)

            # Verify task is cancelled
            status_after = await client_session.experimental.get_task(task_id)
            assert status_after.status == "cancelled"

            tg.cancel_scope.cancel()