"""
Cancellation token implementation for graceful cooperative cancellation.

Provides a thread-safe cancellation primitive that can be shared across
agent execution, tool execution, and streaming operations.
"""

import asyncio
import threading
from concurrent.futures import Future
from typing import Any, Callable, List, Union


class CancellationToken:
    """Thread-safe token for cancelling pending operations."""

    def __init__(self) -> None:
        self._cancelled: bool = False
        self._lock: threading.Lock = threading.Lock()
        self._callbacks: List[Callable[[], None]] = []

    def cancel(self) -> None:
        """Cancel all operations linked to this token."""
        with self._lock:
            if self._cancelled:
                return

            self._cancelled = True
            callbacks = list(self._callbacks)

        for callback in callbacks:
            try:
                callback()
            except Exception:
                pass

    def is_cancelled(self) -> bool:
        """Check if token has been cancelled."""
        with self._lock:
            return self._cancelled

    def add_callback(self, callback: Callable[[], None]) -> None:
        """Register callback invoked when token is cancelled."""
        with self._lock:
            cancelled = self._cancelled
            if not cancelled:
                self._callbacks.append(callback)

        if cancelled:
            try:
                callback()
            except Exception:
                pass

    def link_future(
        self,
        future: Union[Future[Any], asyncio.Future[Any], asyncio.Task[Any]],
    ) -> Union[Future[Any], asyncio.Future[Any], asyncio.Task[Any]]:
        """Link a future/task so it is cancelled when token is cancelled."""
        with self._lock:
            cancelled = self._cancelled
            if not cancelled:

                def _cancel() -> None:
                    future.cancel()

                self._callbacks.append(_cancel)

        if cancelled:
            future.cancel()

        return future


__all__ = ["CancellationToken"]
