"""SEP-1686 client Task classes."""

from __future__ import annotations

import abc
import asyncio
import inspect
import time
import weakref
from collections.abc import Awaitable, Callable
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Generic, TypeVar

import mcp.types
from mcp.types import GetTaskResult, TaskStatusNotification

from fastmcp.client.messages import Message, MessageHandler
from fastmcp.exceptions import ToolError
from fastmcp.utilities.logging import get_logger

logger = get_logger(__name__)

if TYPE_CHECKING:
    from fastmcp.client.client import CallToolResult, Client


class TaskNotificationHandler(MessageHandler):
    """MessageHandler that routes task status notifications to Task objects."""

    def __init__(self, client: Client):
        super().__init__()
        self._client_ref: weakref.ref[Client] = weakref.ref(client)

    async def dispatch(self, message: Message) -> None:
        """Dispatch messages, including task status notifications."""
        if isinstance(message, mcp.types.ServerNotification):
            if isinstance(message.root, TaskStatusNotification):
                client = self._client_ref()
                if client:
                    client._handle_task_status_notification(message.root)

        await super().dispatch(message)


TaskResultT = TypeVar("TaskResultT")


class Task(abc.ABC, Generic[TaskResultT]):
    """
    Abstract base class for MCP background tasks (SEP-1686).

    Provides a uniform API whether the server accepts background execution
    or executes synchronously (graceful degradation per SEP-1686).

    Subclasses:
        - ToolTask: For tool calls (result type: CallToolResult)
        - PromptTask: For prompts (future, result type: GetPromptResult)
        - ResourceTask: For resources (future, result type: ReadResourceResult)
    """

    def __init__(
        self,
        client: Client,
        task_id: str,
        immediate_result: TaskResultT | None = None,
    ):
        """
        Create a Task wrapper.

        Args:
            client: The FastMCP client
            task_id: The task identifier
            immediate_result: If server executed synchronously, the immediate result
        """
        self._client = client
        self._task_id = task_id
        self._immediate_result = immediate_result
        self._is_immediate = immediate_result is not None

        # Notification-based optimization (SEP-1686 notifications/tasks/status)
        self._status_cache: GetTaskResult | None = None
        self._status_event: asyncio.Event | None = None  # Lazy init
        self._status_callbacks: list[
            Callable[[GetTaskResult], None | Awaitable[None]]
        ] = []
        self._cached_result: TaskResultT | None = None

    def _check_client_connected(self) -> None:
        """Validate that client context is still active.

        Raises:
            RuntimeError: If accessed outside client context (unless immediate)
        """
        if self._is_immediate:
            return  # Already resolved, no client needed

        try:
            _ = self._client.session
        except RuntimeError as e:
            raise RuntimeError(
                "Cannot access task results outside client context. "
                "Task futures must be used within 'async with client:' block."
            ) from e

    @property
    def task_id(self) -> str:
        """Get the task ID."""
        return self._task_id

    @property
    def returned_immediately(self) -> bool:
        """Check if server executed the task immediately.

        Returns:
            True if server executed synchronously (graceful degradation or no task support)
            False if server accepted background execution
        """
        return self._is_immediate

    def _handle_status_notification(self, status: GetTaskResult) -> None:
        """Process incoming notifications/tasks/status (internal).

        Called by Client when a notification is received for this task.
        Updates cache, triggers events, and invokes user callbacks.

        Args:
            status: Task status from notification
        """
        # Update cache for next status() call
        self._status_cache = status

        # Wake up any wait() calls
        if self._status_event is not None:
            self._status_event.set()

        # Invoke user callbacks
        for callback in self._status_callbacks:
            try:
                result = callback(status)
                if inspect.isawaitable(result):
                    # Fire and forget async callbacks
                    asyncio.create_task(result)  # type: ignore[arg-type] # noqa: RUF006  # ty:ignore[invalid-argument-type]
            except Exception as e:
                logger.warning(f"Task callback error: {e}", exc_info=True)

    def on_status_change(
        self,
        callback: Callable[[GetTaskResult], None | Awaitable[None]],
    ) -> None:
        """Register callback for status change notifications.

        The callback will be invoked when a notifications/tasks/status is received
        for this task (optional server feature per SEP-1686 lines 436-444).

        Supports both sync and async callbacks (auto-detected).

        Args:
            callback: Function to call with GetTaskResult when status changes.
                     Can return None (sync) or Awaitable[None] (async).

        Example:
            >>> task = await client.call_tool("slow_operation", {}, task=True)
            >>>
            >>> def on_update(status: GetTaskResult):
            ...     print(f"Task {status.taskId} is now {status.status}")
            >>>
            >>> task.on_status_change(on_update)
            >>> result = await task  # Callback fires when status changes
        """
        self._status_callbacks.append(callback)

    async def status(self) -> GetTaskResult:
        """Get current task status.

        If server executed immediately, returns synthetic completed status.
        Otherwise queries the server for current status.
        """
        self._check_client_connected()

        if self._is_immediate:
            # Return synthetic completed status
            now = datetime.now(timezone.utc)
            return GetTaskResult(
                taskId=self._task_id,
                status="completed",
                createdAt=now,
                lastUpdatedAt=now,
                ttl=None,
                pollInterval=1000,
            )

        # Return cached status if available (from notification)
        if self._status_cache is not None:
            cached = self._status_cache
            # Don't clear cache - keep it for next call
            return cached

        # Query server and cache the result
        self._status_cache = await self._client.get_task_status(self._task_id)
        return self._status_cache

    @abc.abstractmethod
    async def result(self) -> TaskResultT:
        """Wait for and return the task result.

        Must be implemented by subclasses to return the appropriate result type.
        """
        ...

    async def wait(
        self, *, state: str | None = None, timeout: float = 300.0
    ) -> GetTaskResult:
        """Wait for task to reach a specific state or complete.

        Uses event-based waiting when notifications are available (fast),
        with fallback to polling (reliable). Optimally wakes up immediately
        on status changes when server sends notifications/tasks/status.

        Args:
            state: Desired state ('working', 'input_required', 'completed', 'failed', 'cancelled').
                   If None, waits until the task exits the 'working' state (completed, failed, cancelled, input_required, etc.)
            timeout: Maximum time to wait in seconds

        Returns:
            GetTaskResult: Final task status

        Raises:
            TimeoutError: If desired state not reached within timeout
        """
        self._check_client_connected()

        if self._is_immediate:
            # Already done
            return await self.status()

        # Initialize event for notification wake-ups
        if self._status_event is None:
            self._status_event = asyncio.Event()

        start = time.time()
        in_progress_states = {"working"}
        poll_interval = 0.5  # Fallback polling interval (500ms)

        while True:
            # Check cached status first (updated by notifications)
            if self._status_cache:
                current = self._status_cache.status
                if state is None:
                    if current not in in_progress_states:
                        return self._status_cache
                elif current == state:
                    return self._status_cache

            # Check timeout
            elapsed = time.time() - start
            if elapsed >= timeout:
                raise TimeoutError(
                    f"Task {self._task_id} did not reach {state or 'terminal state'} within {timeout}s"
                )

            remaining = timeout - elapsed

            # Wait for notification event OR poll timeout
            try:
                await asyncio.wait_for(
                    self._status_event.wait(), timeout=min(poll_interval, remaining)
                )
                self._status_event.clear()
            except asyncio.TimeoutError:
                # Fallback: poll server (notification didn't arrive in time)
                self._status_cache = await self._client.get_task_status(self._task_id)

    async def _wait_terminal(self, timeout: float = 300.0) -> GetTaskResult:
        """Wait until task reaches a terminal state (completed, failed, cancelled).

        Unlike wait(), this will not return on input_required — it continues
        waiting until the task fully resolves. Used internally by result().
        """
        terminal_states = {"completed", "failed", "cancelled"}
        status = await self.wait(timeout=timeout)
        while status.status not in terminal_states:
            # Task is in a non-terminal state (e.g. input_required) — reset
            # cache so the next wait() call blocks instead of returning immediately.
            self._status_cache = None
            status = await self.wait(timeout=timeout)
        return status

    async def cancel(self) -> None:
        """Cancel this task, transitioning it to cancelled state.

        Sends a tasks/cancel protocol request. The server will attempt to halt
        execution and move the task to cancelled state.

        Note: If server executed immediately (graceful degradation), this is a no-op
        as there's no server-side task to cancel.
        """
        if self._is_immediate:
            # No server-side task to cancel
            return
        self._check_client_connected()
        await self._client.cancel_task(self._task_id)
        # Invalidate cache to force fresh status fetch
        self._status_cache = None

    def __await__(self):
        """Allow 'await task' to get result."""
        return self.result().__await__()


class ToolTask(Task["CallToolResult"]):
    """
    Represents a tool call that may execute in background or immediately.

    Provides a uniform API whether the server accepts background execution
    or executes synchronously (graceful degradation per SEP-1686).

    Usage:
        task = await client.call_tool_as_task("analyze", args)

        # Check status
        status = await task.status()

        # Wait for completion
        await task.wait()

        # Get result (waits if needed)
        result = await task.result()  # Returns CallToolResult

        # Or just await the task directly
        result = await task
    """

    def __init__(
        self,
        client: Client,
        task_id: str,
        tool_name: str,
        immediate_result: CallToolResult | None = None,
        raise_on_error: bool = True,
    ):
        """
        Create a ToolTask wrapper.

        Args:
            client: The FastMCP client
            task_id: The task identifier
            tool_name: Name of the tool being executed
            immediate_result: If server executed synchronously, the immediate result
            raise_on_error: Whether task.result() should raise ToolError on errors
        """
        super().__init__(client, task_id, immediate_result)
        self._tool_name = tool_name
        self._raise_on_error = raise_on_error

    async def result(self) -> CallToolResult:
        """Wait for and return the tool result.

        If server executed immediately, returns the immediate result.
        Otherwise waits for background task to complete and retrieves result.

        Returns:
            CallToolResult: The parsed tool result (same as call_tool returns)
        """
        # Check cache first
        if self._cached_result is not None:
            return self._cached_result

        if self._is_immediate:
            assert self._immediate_result is not None  # Type narrowing
            result = self._immediate_result
            if result.is_error and self._raise_on_error:
                if result.content and isinstance(
                    result.content[0], mcp.types.TextContent
                ):
                    msg = result.content[0].text
                else:
                    msg = f"Tool '{self._tool_name}' returned an error"
                raise ToolError(msg)
        else:
            # Check client connected
            self._check_client_connected()

            # Wait for completion using event-based wait (respects notifications)
            await self._wait_terminal()

            # Get the raw result (dict or CallToolResult)
            raw_result = await self._client.get_task_result(self._task_id)

            # Convert to CallToolResult if needed and parse
            if isinstance(raw_result, dict):
                # Raw dict from get_task_result - parse as CallToolResult
                mcp_result = mcp.types.CallToolResult.model_validate(raw_result)
                result = await self._client._parse_call_tool_result(
                    self._tool_name,
                    mcp_result,
                    raise_on_error=self._raise_on_error,
                )
            elif isinstance(raw_result, mcp.types.CallToolResult):
                # Already a CallToolResult from MCP protocol - parse it
                result = await self._client._parse_call_tool_result(
                    self._tool_name,
                    raw_result,
                    raise_on_error=self._raise_on_error,
                )
            else:
                # Legacy ToolResult format - convert to MCP type
                if hasattr(raw_result, "content") and hasattr(
                    raw_result, "structured_content"
                ):
                    mcp_result = mcp.types.CallToolResult(
                        content=raw_result.content,
                        structuredContent=raw_result.structured_content,
                        _meta=raw_result.meta,  # type: ignore[call-arg]  # _meta is Pydantic alias for meta field
                    )
                    result = await self._client._parse_call_tool_result(
                        self._tool_name,
                        mcp_result,
                        raise_on_error=self._raise_on_error,
                    )
                else:
                    # Unknown type - just return it
                    result = raw_result

        # Cache before returning
        self._cached_result = result
        return result


class PromptTask(Task[mcp.types.GetPromptResult]):
    """
    Represents a prompt call that may execute in background or immediately.

    Provides a uniform API whether the server accepts background execution
    or executes synchronously (graceful degradation per SEP-1686).

    Usage:
        task = await client.get_prompt_as_task("analyze", args)
        result = await task  # Returns GetPromptResult
    """

    def __init__(
        self,
        client: Client,
        task_id: str,
        prompt_name: str,
        immediate_result: mcp.types.GetPromptResult | None = None,
    ):
        """
        Create a PromptTask wrapper.

        Args:
            client: The FastMCP client
            task_id: The task identifier
            prompt_name: Name of the prompt being executed
            immediate_result: If server executed synchronously, the immediate result
        """
        super().__init__(client, task_id, immediate_result)
        self._prompt_name = prompt_name

    async def result(self) -> mcp.types.GetPromptResult:
        """Wait for and return the prompt result.

        If server executed immediately, returns the immediate result.
        Otherwise waits for background task to complete and retrieves result.

        Returns:
            GetPromptResult: The prompt result with messages and description
        """
        # Check cache first
        if self._cached_result is not None:
            return self._cached_result

        if self._is_immediate:
            assert self._immediate_result is not None
            result = self._immediate_result
        else:
            # Check client connected
            self._check_client_connected()

            # Wait for completion using event-based wait (respects notifications)
            await self._wait_terminal()

            # Get the raw MCP result
            mcp_result = await self._client.get_task_result(self._task_id)

            # Parse as GetPromptResult
            result = mcp.types.GetPromptResult.model_validate(mcp_result)

        # Cache before returning
        self._cached_result = result
        return result


class ResourceTask(
    Task[list[mcp.types.TextResourceContents | mcp.types.BlobResourceContents]]
):
    """
    Represents a resource read that may execute in background or immediately.

    Provides a uniform API whether the server accepts background execution
    or executes synchronously (graceful degradation per SEP-1686).

    Usage:
        task = await client.read_resource_as_task("file://data.txt")
        contents = await task  # Returns list[ReadResourceContents]
    """

    def __init__(
        self,
        client: Client,
        task_id: str,
        uri: str,
        immediate_result: list[
            mcp.types.TextResourceContents | mcp.types.BlobResourceContents
        ]
        | None = None,
    ):
        """
        Create a ResourceTask wrapper.

        Args:
            client: The FastMCP client
            task_id: The task identifier
            uri: URI of the resource being read
            immediate_result: If server executed synchronously, the immediate result
        """
        super().__init__(client, task_id, immediate_result)
        self._uri = uri

    async def result(
        self,
    ) -> list[mcp.types.TextResourceContents | mcp.types.BlobResourceContents]:
        """Wait for and return the resource contents.

        If server executed immediately, returns the immediate result.
        Otherwise waits for background task to complete and retrieves result.

        Returns:
            list[ReadResourceContents]: The resource contents
        """
        # Check cache first
        if self._cached_result is not None:
            return self._cached_result

        if self._is_immediate:
            assert self._immediate_result is not None
            result = self._immediate_result
        else:
            # Check client connected
            self._check_client_connected()

            # Wait for completion using event-based wait (respects notifications)
            await self._wait_terminal()

            # Get the raw MCP result
            mcp_result = await self._client.get_task_result(self._task_id)

            # Parse as ReadResourceResult or extract contents
            if isinstance(mcp_result, mcp.types.ReadResourceResult):
                # Already parsed by TasksResponse - extract contents
                result = list(mcp_result.contents)
            elif isinstance(mcp_result, dict) and "contents" in mcp_result:
                # Dict format - parse each content item
                parsed_contents = []
                for item in mcp_result["contents"]:
                    if isinstance(item, dict):
                        if "blob" in item:
                            parsed_contents.append(
                                mcp.types.BlobResourceContents.model_validate(item)
                            )
                        else:
                            parsed_contents.append(
                                mcp.types.TextResourceContents.model_validate(item)
                            )
                    else:
                        parsed_contents.append(item)
                result = parsed_contents
            else:
                # Fallback - might be the list directly
                result = mcp_result if isinstance(mcp_result, list) else [mcp_result]

        # Cache before returning
        self._cached_result = result
        return result