"""
Cancellation framework based on CancellationToken + Scope pattern.

Key design choices:
- **Sticky cancellation**: once a token is cancelled, it stays cancelled until
  the owning scope exits (or an explicit reset). Multiple observers can all
  see the cancelled state.
- **CancellationScope**: context-manager that auto-registers / unregisters a
  token and propagates it via ``contextvars`` so deep call-chains don't need
  to pass the token explicitly.
- **Thread-safe**: all mutable state is guarded by a single ``threading.Lock``
  and iterations always operate on snapshots.
- **Backward-compatible**: the existing ``global_cancel`` singleton and every
  public method name is preserved.  New capabilities are additive.
"""

import contextvars
import threading
import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Generator, List, Optional, Set

# ---------------------------------------------------------------------------
# ContextVar – carries the "current cancel token" across sync call-chains
# and is automatically copied into child threads started via
# ``threading.Thread`` (Python 3.12+) or ``concurrent.futures``.
# ---------------------------------------------------------------------------
_current_cancel_token: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    "current_cancel_token", default=None
)


# ---------------------------------------------------------------------------
# Exception
# ---------------------------------------------------------------------------
class CancelRequestedException(Exception):
    """当取消请求被触发时抛出的异常"""

    def __init__(
        self, token: Optional[str] = None, message: str = "Operation was cancelled"
    ):
        self.token = token
        self.message = message
        super().__init__(self.message)


# ---------------------------------------------------------------------------
# Internal per-token state
# ---------------------------------------------------------------------------
@dataclass
class _TokenState:
    """Internal mutable state for a single cancellation token."""

    cancelled: bool = False
    active_count: int = 0  # supports nested scopes for the same token
    context: Dict[str, Any] = field(default_factory=dict)
    created_at: float = field(default_factory=time.time)
    callbacks: List[Callable[[], None]] = field(default_factory=list)


# ---------------------------------------------------------------------------
# GlobalCancel – the singleton registry
# ---------------------------------------------------------------------------
class GlobalCancel:
    """Thread-safe cancellation registry with scope support.

    All pre-existing public methods are preserved for backward compatibility.
    New capabilities:

    * ``scope(token)`` – context-manager for auto register/unregister.
    * ``raise_if_cancelled(token)`` – sticky check (does **not** consume the
      cancellation state).
    * ``cancel(token, context)`` – alias for ``set(token, context)``; also
      fires registered callbacks.
    * ``cancel_active(context)`` – thread-safe batch cancel of all active
      tokens (snapshot-based, no iteration-mutation risk).
    * ``on_cancel(token, callback)`` – register a callback invoked when the
      token is cancelled.
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._global_flag: bool = False
        self._global_context: Dict[str, Any] = {}
        self._tokens: Dict[str, _TokenState] = {}
        self._active_tokens: Set[str] = set()
        # kept for legacy callers that read it directly
        self._current_task_token: Optional[str] = None

    # ------------------------------------------------------------------
    # helpers
    # ------------------------------------------------------------------
    def _get_or_create_state(self, token: str) -> _TokenState:
        """Return existing state or lazily create one (caller holds lock)."""
        if token not in self._tokens:
            self._tokens[token] = _TokenState()
        return self._tokens[token]

    def _effective_token(self, token: Optional[str]) -> Optional[str]:
        """Resolve *token*: if ``None`` fall back to the contextvar."""
        if token is not None:
            return token
        return _current_cancel_token.get(None)

    # ------------------------------------------------------------------
    # Scope (new API)
    # ------------------------------------------------------------------
    @contextmanager
    def scope(
        self, token: Optional[str] = None
    ) -> Generator[Optional[str], None, None]:
        """Context-manager that registers *token* on entry and unregisters on
        exit.  While the scope is active the contextvar
        ``_current_cancel_token`` is set so that ``check_and_raise()`` /
        ``raise_if_cancelled()`` can discover it automatically.

        Usage::

            with global_cancel.scope(cancel_token):
                do_work()          # check_and_raise() will use cancel_token
                call_deep_fn()     # same, no need to pass token explicitly
        """
        if token is not None:
            self.register_token(token)
        reset_cv = _current_cancel_token.set(token)
        try:
            yield token
        finally:
            _current_cancel_token.reset(reset_cv)
            if token is not None:
                self._unregister_token(token)

    # ------------------------------------------------------------------
    # Registration
    # ------------------------------------------------------------------
    def register_token(self, token: str) -> None:
        """注册一个 token，表示一个操作开始，但尚未请求取消

        If the token was already cancelled before registration (e.g. a race
        where ``set()`` arrived first) the cancelled state is **preserved**,
        not overwritten.
        """
        with self._lock:
            state = self._get_or_create_state(token)
            state.active_count += 1
            self._active_tokens.add(token)
            self._current_task_token = token

    def _unregister_token(self, token: str) -> None:
        """Decrement active count; remove state when no scopes remain."""
        with self._lock:
            state = self._tokens.get(token)
            if state is None:
                return
            state.active_count = max(0, state.active_count - 1)
            if state.active_count <= 0:
                self._active_tokens.discard(token)
                del self._tokens[token]
                if self._current_task_token == token:
                    self._current_task_token = None

    # ------------------------------------------------------------------
    # Query
    # ------------------------------------------------------------------
    def get_current_task_token(self) -> Optional[str]:
        """获取当前任务token"""
        with self._lock:
            return self._current_task_token

    def get_active_tokens(self) -> Set[str]:
        """获取当前正在运行的token"""
        with self._lock:
            return self._active_tokens.copy()

    def is_requested(self, token: Optional[str] = None) -> bool:
        """检查是否请求了特定token或全局的取消"""
        with self._lock:
            if self._global_flag:
                return True
            resolved = self._effective_token(token)
            if resolved is not None:
                state = self._tokens.get(resolved)
                if state is not None:
                    return state.cancelled
            return False

    # ------------------------------------------------------------------
    # Cancel (set)
    # ------------------------------------------------------------------
    def set(
        self,
        token: Optional[str] = None,
        context: Optional[Dict[str, Any]] = None,
    ) -> None:
        """设置特定token或全局的取消标志"""
        callbacks_to_fire: List[Callable[[], None]] = []
        with self._lock:
            if token is None:
                self._global_flag = True
                if context:
                    self._global_context.update(context)
            else:
                state = self._get_or_create_state(token)
                if not state.cancelled:
                    state.cancelled = True
                    callbacks_to_fire = list(state.callbacks)
                if context:
                    state.context.update(context)

        # Fire callbacks outside the lock to avoid deadlocks
        for cb in callbacks_to_fire:
            try:
                cb()
            except Exception:
                pass

    # Alias – more expressive name
    cancel = set

    def set_active_tokens(self) -> None:
        """启用所有活跃的token (thread-safe snapshot iteration)."""
        with self._lock:
            snapshot = list(self._active_tokens)
        for t in snapshot:
            self.set(t)

    # Alias
    cancel_active = set_active_tokens

    # ------------------------------------------------------------------
    # Reset
    # ------------------------------------------------------------------
    def reset_global(self) -> None:
        """重置全局取消标志"""
        with self._lock:
            self._global_flag = False
            self._global_context.clear()

    def reset_token(self, token: str) -> None:
        """重置特定token的取消标志并从活跃集合移除。"""
        with self._lock:
            if token in self._tokens:
                del self._tokens[token]
            self._active_tokens.discard(token)
            if self._current_task_token == token:
                self._current_task_token = None

    def reset(self, token: Optional[str] = None) -> None:
        """重置特定token或全局的取消标志"""
        with self._lock:
            if token is None:
                self._global_flag = False
                self._global_context.clear()
                self._tokens.clear()
                self._active_tokens.clear()
                self._current_task_token = None
            else:
                if token in self._tokens:
                    del self._tokens[token]
                self._active_tokens.discard(token)
                if self._current_task_token == token:
                    self._current_task_token = None

    def reset_active_tokens(self) -> None:
        """重置所有活跃的token"""
        with self._lock:
            tokens_snapshot = list(self._active_tokens)
        for t in tokens_snapshot:
            self.reset_token(t)

    # ------------------------------------------------------------------
    # Context
    # ------------------------------------------------------------------
    def get_context(self, token: Optional[str] = None) -> Dict[str, Any]:
        """获取与取消相关的上下文信息"""
        with self._lock:
            if token is None:
                return self._global_context.copy()
            state = self._tokens.get(token)
            if state is not None:
                return state.context.copy()
            return {}

    # ------------------------------------------------------------------
    # Check & raise helpers
    # ------------------------------------------------------------------
    def raise_if_cancelled(self, token: Optional[str] = None) -> None:
        """Sticky check: raises ``CancelRequestedException`` if the token
        (or the global flag) is cancelled.  Does **not** reset/consume the
        cancellation state – that is left to the owning ``scope``.

        If *token* is ``None`` the current contextvar token is used.
        """
        resolved = self._effective_token(token)
        if self.is_requested(resolved):
            ctx = self.get_context(resolved)
            raise CancelRequestedException(
                resolved, ctx.get("message", "Operation was cancelled")
            )

    def check_and_raise(self, token: Optional[str] = None) -> None:
        """检查是否请求了取消，如果是则抛出异常

        **Backward-compatible**: this method now delegates to
        ``raise_if_cancelled`` (sticky semantics).  The cancellation state is
        **not** consumed – cleanup happens when the scope exits or when
        ``reset_token`` / ``reset`` is called explicitly.
        """
        self.raise_if_cancelled(token)

    # ------------------------------------------------------------------
    # Callback registration (new API)
    # ------------------------------------------------------------------
    def on_cancel(self, token: str, callback: Callable[[], None]) -> None:
        """Register a callback to be invoked when *token* is cancelled.

        If *token* is **already** cancelled at the time of registration the
        callback is fired immediately (outside the lock).
        """
        fire_now = False
        with self._lock:
            state = self._get_or_create_state(token)
            if state.cancelled:
                fire_now = True
            else:
                state.callbacks.append(callback)
        if fire_now:
            try:
                callback()
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Legacy property kept for code that accesses _token_flags directly
    # ------------------------------------------------------------------
    @property
    def _token_flags(self) -> Dict[str, bool]:
        """Backward-compat shim: returns a snapshot dict mirroring old
        ``_token_flags`` layout.  **Read-only** – mutations are ignored.
        """
        with self._lock:
            return {t: s.cancelled for t, s in self._tokens.items()}


global_cancel = GlobalCancel()
