"""Experimental handlers for the low-level MCP server.

WARNING: These APIs are experimental and may change without notice.
"""

from __future__ import annotations

import logging
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING

from mcp.server.experimental.task_scope import task_in_session_scope, task_listable_in_session_scope
from mcp.server.experimental.task_support import TaskSupport
from mcp.server.lowlevel.func_inspection import create_call_wrapper
from mcp.shared.exceptions import McpError
from mcp.shared.experimental.tasks.helpers import cancel_task
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue
from mcp.shared.experimental.tasks.store import TaskStore
from mcp.types import (
    INVALID_PARAMS,
    CancelTaskRequest,
    CancelTaskResult,
    ErrorData,
    GetTaskPayloadRequest,
    GetTaskPayloadResult,
    GetTaskRequest,
    GetTaskResult,
    ListTasksRequest,
    ListTasksResult,
    ServerCapabilities,
    ServerResult,
    ServerTasksCapability,
    ServerTasksRequestsCapability,
    Task,
    TasksCallCapability,
    TasksCancelCapability,
    TasksListCapability,
    TasksToolsCapability,
)

if TYPE_CHECKING:
    from mcp.server.lowlevel.server import Server

logger = logging.getLogger(__name__)


class ExperimentalHandlers:
    """Experimental request/notification handlers.

    WARNING: These APIs are experimental and may change without notice.
    """

    def __init__(
        self,
        server: Server,
        request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]],
        notification_handlers: dict[type, Callable[..., Awaitable[None]]],
    ):
        self._server = server
        self._request_handlers = request_handlers
        self._notification_handlers = notification_handlers
        self._task_support: TaskSupport | None = None

    @property
    def task_support(self) -> TaskSupport | None:
        """Get the task support configuration, if enabled."""
        return self._task_support

    def update_capabilities(self, capabilities: ServerCapabilities) -> None:
        # Only add tasks capability if handlers are registered
        if not any(
            req_type in self._request_handlers
            for req_type in [GetTaskRequest, ListTasksRequest, CancelTaskRequest, GetTaskPayloadRequest]
        ):
            return

        capabilities.tasks = ServerTasksCapability()
        if ListTasksRequest in self._request_handlers:
            capabilities.tasks.list = TasksListCapability()
        if CancelTaskRequest in self._request_handlers:
            capabilities.tasks.cancel = TasksCancelCapability()

        capabilities.tasks.requests = ServerTasksRequestsCapability(
            tools=TasksToolsCapability(call=TasksCallCapability())
        )  # assuming always supported for now

    def enable_tasks(
        self,
        store: TaskStore | None = None,
        queue: TaskMessageQueue | None = None,
    ) -> TaskSupport:
        """
        Enable experimental task support.

        This sets up the task infrastructure and auto-registers default handlers
        for tasks/get, tasks/result, tasks/list, and tasks/cancel.

        Args:
            store: Custom TaskStore implementation (defaults to InMemoryTaskStore)
            queue: Custom TaskMessageQueue implementation (defaults to InMemoryTaskMessageQueue)

        Returns:
            The TaskSupport configuration object

        Example:
            # Simple in-memory setup
            server.experimental.enable_tasks()

            # Custom store/queue for distributed systems
            server.experimental.enable_tasks(
                store=RedisTaskStore(redis_url),
                queue=RedisTaskMessageQueue(redis_url),
            )

        WARNING: This API is experimental and may change without notice.
        """
        if store is None:
            store = InMemoryTaskStore()
        if queue is None:
            queue = InMemoryTaskMessageQueue()

        self._task_support = TaskSupport(store=store, queue=queue)

        # Auto-register default handlers
        self._register_default_task_handlers()

        return self._task_support

    def _requestor_session_scope(self) -> str | None:
        """Return the task session scope of the session making the current request."""
        session = self._server.request_context.session
        return session._experimental.task_session_scope  # pyright: ignore[reportPrivateUsage]

    def _require_task_in_requestor_scope(self, task_id: str) -> None:
        """Reject task IDs that belong to a different session.

        Task IDs generated by `run_task()` embed the creating session's
        scope. The default handlers treat a task created by another session
        exactly like a task that does not exist, so a requestor cannot tell
        whether such a task exists. Task IDs without an embedded scope are
        accepted from any session.

        Raises:
            McpError: With INVALID_PARAMS if the task belongs to another session.
        """
        if not task_in_session_scope(task_id, self._requestor_session_scope()):
            raise McpError(
                ErrorData(
                    code=INVALID_PARAMS,
                    message=f"Task not found: {task_id}",
                )
            )

    def _register_default_task_handlers(self) -> None:
        """Register default handlers for task operations.

        Each default handler only operates on tasks created by the requesting
        session (see `_require_task_in_requestor_scope`), and tasks/list only
        returns the requesting session's own tasks (see
        `task_listable_in_session_scope`).
        """
        assert self._task_support is not None
        support = self._task_support

        # Register get_task handler if not already registered
        if GetTaskRequest not in self._request_handlers:

            async def _default_get_task(req: GetTaskRequest) -> ServerResult:
                self._require_task_in_requestor_scope(req.params.taskId)
                task = await support.store.get_task(req.params.taskId)
                if task is None:
                    raise McpError(
                        ErrorData(
                            code=INVALID_PARAMS,
                            message=f"Task not found: {req.params.taskId}",
                        )
                    )
                return ServerResult(
                    GetTaskResult(
                        taskId=task.taskId,
                        status=task.status,
                        statusMessage=task.statusMessage,
                        createdAt=task.createdAt,
                        lastUpdatedAt=task.lastUpdatedAt,
                        ttl=task.ttl,
                        pollInterval=task.pollInterval,
                    )
                )

            self._request_handlers[GetTaskRequest] = _default_get_task

        # Register get_task_result handler if not already registered
        if GetTaskPayloadRequest not in self._request_handlers:

            async def _default_get_task_result(req: GetTaskPayloadRequest) -> ServerResult:
                self._require_task_in_requestor_scope(req.params.taskId)
                ctx = self._server.request_context
                result = await support.handler.handle(req, ctx.session, ctx.request_id)
                return ServerResult(result)

            self._request_handlers[GetTaskPayloadRequest] = _default_get_task_result

        # Register list_tasks handler if not already registered
        if ListTasksRequest not in self._request_handlers:

            async def _default_list_tasks(req: ListTasksRequest) -> ServerResult:
                requestor_scope = self._requestor_session_scope()
                if requestor_scope is None:
                    # The server cannot tell this requestor apart from any
                    # other, so there are no tasks it can be shown.
                    return ServerResult(ListTasksResult(tasks=[]))
                # Return every task that belongs to the requesting session in
                # a single page. The store's pagination cursor is never sent
                # to the requestor: it is derived from the unfiltered listing,
                # so it could identify a task belonging to a different
                # session. For the same reason the request's cursor is not
                # forwarded to the store.
                own_tasks: list[Task] = []
                cursor: str | None = None
                while True:
                    page, cursor = await support.store.list_tasks(cursor)
                    own_tasks.extend(
                        task for task in page if task_listable_in_session_scope(task.taskId, requestor_scope)
                    )
                    if cursor is None:
                        return ServerResult(ListTasksResult(tasks=own_tasks))

            self._request_handlers[ListTasksRequest] = _default_list_tasks

        # Register cancel_task handler if not already registered
        if CancelTaskRequest not in self._request_handlers:

            async def _default_cancel_task(req: CancelTaskRequest) -> ServerResult:
                self._require_task_in_requestor_scope(req.params.taskId)
                result = await cancel_task(support.store, req.params.taskId)
                return ServerResult(result)

            self._request_handlers[CancelTaskRequest] = _default_cancel_task

    def list_tasks(
        self,
    ) -> Callable[
        [Callable[[ListTasksRequest], Awaitable[ListTasksResult]]],
        Callable[[ListTasksRequest], Awaitable[ListTasksResult]],
    ]:
        """Register a handler for listing tasks.

        WARNING: This API is experimental and may change without notice.
        """

        def decorator(
            func: Callable[[ListTasksRequest], Awaitable[ListTasksResult]],
        ) -> Callable[[ListTasksRequest], Awaitable[ListTasksResult]]:
            logger.debug("Registering handler for ListTasksRequest")
            wrapper = create_call_wrapper(func, ListTasksRequest)

            async def handler(req: ListTasksRequest) -> ServerResult:
                result = await wrapper(req)
                return ServerResult(result)

            self._request_handlers[ListTasksRequest] = handler
            return func

        return decorator

    def get_task(
        self,
    ) -> Callable[
        [Callable[[GetTaskRequest], Awaitable[GetTaskResult]]], Callable[[GetTaskRequest], Awaitable[GetTaskResult]]
    ]:
        """Register a handler for getting task status.

        WARNING: This API is experimental and may change without notice.
        """

        def decorator(
            func: Callable[[GetTaskRequest], Awaitable[GetTaskResult]],
        ) -> Callable[[GetTaskRequest], Awaitable[GetTaskResult]]:
            logger.debug("Registering handler for GetTaskRequest")
            wrapper = create_call_wrapper(func, GetTaskRequest)

            async def handler(req: GetTaskRequest) -> ServerResult:
                result = await wrapper(req)
                return ServerResult(result)

            self._request_handlers[GetTaskRequest] = handler
            return func

        return decorator

    def get_task_result(
        self,
    ) -> Callable[
        [Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]],
        Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]],
    ]:
        """Register a handler for getting task results/payload.

        WARNING: This API is experimental and may change without notice.
        """

        def decorator(
            func: Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]],
        ) -> Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]:
            logger.debug("Registering handler for GetTaskPayloadRequest")
            wrapper = create_call_wrapper(func, GetTaskPayloadRequest)

            async def handler(req: GetTaskPayloadRequest) -> ServerResult:
                result = await wrapper(req)
                return ServerResult(result)

            self._request_handlers[GetTaskPayloadRequest] = handler
            return func

        return decorator

    def cancel_task(
        self,
    ) -> Callable[
        [Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]],
        Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]],
    ]:
        """Register a handler for cancelling tasks.

        WARNING: This API is experimental and may change without notice.
        """

        def decorator(
            func: Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]],
        ) -> Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]:
            logger.debug("Registering handler for CancelTaskRequest")
            wrapper = create_call_wrapper(func, CancelTaskRequest)

            async def handler(req: CancelTaskRequest) -> ServerResult:
                result = await wrapper(req)
                return ServerResult(result)

            self._request_handlers[CancelTaskRequest] = handler
            return func

        return decorator