from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Generic, Optional

from nexusrpc._common import Link, OutputT


class OperationTaskCancellation(ABC):
    """
    Indicates whether a a Nexus task has been cancelled during a sync operation or before an async operation has
    returned a token.

    Nexus worker implementations are expected to provide an implementation that enables
    cooperative cancellation for both sync and async operation handlers.

    Operation Handler implementations are expected to periodically check :py:attr:`is_cancelled` or use :py:attr:`wait_until_cancelled` / :py:attr:`wait_until_cancelled_sync` to cancel in flight work when appropriate.
    """

    @abstractmethod
    def is_cancelled(self) -> bool:
        """
        Return True if the associated task has been cancelled.
        """
        raise NotImplementedError

    @abstractmethod
    def cancellation_reason(self) -> Optional[str]:
        """
        Provide additional context for the cancellation, if available.
        """
        raise NotImplementedError

    @abstractmethod
    def wait_until_cancelled_sync(self, timeout: Optional[float] = None) -> bool:
        """
        Block until cancellation occurs or the optional timeout elapses. Nexus worker implementations may return `True` for :py:attr:`is_cancelled` before this method returns and therefore may cause a race condition if both are used in tandem.
        """
        raise NotImplementedError

    @abstractmethod
    async def wait_until_cancelled(self) -> None:
        """Await cancellation using async primitives. Nexus worker implementations may return `True` for :py:attr:`is_cancelled` before this method returns and therefore may cause a race condition if both are used in tandem."""
        raise NotImplementedError


@dataclass(frozen=True, kw_only=True)
class OperationContext(ABC):
    """
    Context for the execution of the requested operation method.

    Includes information from the request.
    """

    def __new__(cls, *args: Any, **kwargs: Any):
        if cls is OperationContext:
            raise TypeError(
                "OperationContext is an abstract class and cannot be instantiated directly"
            )
        return super().__new__(cls)

    service: str
    """
    The name of the service that the operation belongs to.
    """

    operation: str
    """
    The name of the operation.
    """

    headers: Mapping[str, str]
    """
    Optional header fields sent by the caller.
    """

    task_cancellation: OperationTaskCancellation
    """
    Task cancellation information indicating that a running task should be interrupted. This is distinct from operation cancellation.
    """

    request_deadline: Optional[datetime] = None
    """
    The deadline for the operation handler method. Note that this is the time by which the
    current _request_ should complete, not the _operation_'s deadline.
    """


@dataclass(frozen=True, kw_only=True)
class StartOperationContext(OperationContext):
    """
    Context for the start method.

    Includes information from the request.
    """

    request_id: str
    """
    Request ID that may be used by the handler to dedupe a start request.
    By default a v4 UUID should be generated by the client.
    """

    callback_url: Optional[str] = None
    """
    A callback URL is required to deliver the completion of an async operation. This URL should be
    called by a handler upon completion if the started operation is async.
    """

    callback_headers: Mapping[str, str] = field(default_factory=dict)
    """
    Optional header fields set by the caller to be attached to the callback request when an
    asynchronous operation completes.
    """

    inbound_links: Sequence[Link] = field(default_factory=list)
    """
    Links received in the request. This list is automatically populated when handling a start
    request. Handlers may use these links, for example to add information about the
    caller to a resource associated with the operation execution.
    """

    outbound_links: list[Link] = field(default_factory=list)
    """
    Links to be returned by the handler. This list is initially empty. Handlers may
    populate this list, for example with links describing resources associated with the
    operation execution that may be of use to the caller.
    """


@dataclass(frozen=True, kw_only=True)
class CancelOperationContext(OperationContext):
    """
    Context for the cancel method.

    Includes information from the request.
    """


@dataclass(frozen=True)
class StartOperationResultSync(Generic[OutputT]):
    """
    A result returned synchronously by the start method of a nexus operation handler.
    """

    value: OutputT  # type: ignore[misc]
    """
    The value returned by the operation.
    """


@dataclass(frozen=True)
class StartOperationResultAsync:
    """
    A value returned by the start method of a nexus operation handler indicating that
    the operation is responding asynchronously.
    """

    token: str
    """
    A token representing the in-progress operation that the caller can submit with
    subsequent ``fetch_info``, ``fetch_result``, or ``cancel`` requests.
    """
