"""Cancellation token for aborting async operations."""

import asyncio
from dataclasses import dataclass, field
from typing import Callable


@dataclass
class CancellationToken:
    """Token for cooperative cancellation of async operations.

    Provides a thread-safe way to signal cancellation to running tasks.
    Tasks should check is_cancelled periodically and stop gracefully.

    Usage:
        token = CancellationToken()

        async def long_task():
            while not token.is_cancelled:
                await do_work()
            # Clean up

        # Later, to cancel:
        token.cancel()

    With callbacks:
        token = CancellationToken()
        token.on_cancel(cleanup_function)
        token.cancel()  # cleanup_function is called
    """

    _cancelled: bool = field(default=False, init=False)
    _callbacks: list[Callable[[], None]] = field(default_factory=list, init=False)
    _cancel_event: asyncio.Event | None = field(default=None, init=False)

    @property
    def is_cancelled(self) -> bool:
        """Whether cancellation has been requested."""
        return self._cancelled

    def cancel(self) -> None:
        """Request cancellation.

        This is idempotent - calling multiple times has no additional effect.
        """
        if self._cancelled:
            return

        self._cancelled = True

        # Notify callbacks
        for callback in self._callbacks:
            try:
                callback()
            except Exception:
                pass  # Don't let callback errors propagate

        # Set event for async waiters
        if self._cancel_event is not None:
            self._cancel_event.set()

    def on_cancel(self, callback: Callable[[], None]) -> None:
        """Register a callback to be invoked on cancellation.

        If already cancelled, callback is invoked immediately.

        Args:
            callback: Function to call when cancelled.
        """
        if self._cancelled:
            try:
                callback()
            except Exception:
                pass
        else:
            self._callbacks.append(callback)

    def reset(self) -> None:
        """Reset the token for reuse.

        Clears cancelled state and callbacks.
        """
        self._cancelled = False
        self._callbacks.clear()
        if self._cancel_event is not None:
            self._cancel_event.clear()

    async def wait(self, timeout: float | None = None) -> bool:
        """Wait until cancellation is requested.

        Args:
            timeout: Maximum time to wait in seconds.

        Returns:
            True if cancelled, False if timed out.
        """
        if self._cancelled:
            return True

        if self._cancel_event is None:
            self._cancel_event = asyncio.Event()

        try:
            await asyncio.wait_for(self._cancel_event.wait(), timeout)
            return True
        except asyncio.TimeoutError:
            return False

    def raise_if_cancelled(self) -> None:
        """Raise CancellationError if cancelled.

        Use this for tasks that should stop immediately on cancel.
        """
        if self._cancelled:
            raise CancellationError("Operation was cancelled")


class CancellationError(Exception):
    """Raised when an operation is cancelled."""

    pass


class CancellationScope:
    """Context manager for cancellation tokens.

    Creates a token that can be cancelled from outside the scope.

    Usage:
        scope = CancellationScope()
        with scope:
            while not scope.token.is_cancelled:
                await do_work()

        # From elsewhere:
        scope.cancel()
    """

    def __init__(self) -> None:
        self._token = CancellationToken()

    @property
    def token(self) -> CancellationToken:
        """The cancellation token for this scope."""
        return self._token

    @property
    def is_cancelled(self) -> bool:
        """Whether the scope has been cancelled."""
        return self._token.is_cancelled

    def cancel(self) -> None:
        """Cancel the scope."""
        self._token.cancel()

    def __enter__(self) -> "CancellationScope":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        # Auto-cancel on exit if not already cancelled
        if not self._token.is_cancelled:
            self._token.cancel()
