"""Cooperative cancellation token for deep-research jobs.

The vendored GPT Researcher loops (recursive deep research, sub-topic fan-out,
multi-agent review cycles) can run for minutes.  Every adapter and runner polls
a single :class:`CancellationToken` so a user cancel (or a lease takeover)
promptly unwinds the whole tree instead of burning model budget to completion.

Design:

- ``CancellationToken`` is an asyncio-friendly flag — ``cancel()`` flips it,
  ``raise_if_cancelled()`` is the checked-call sites sprinkle at loop boundaries.
- A token is created **per job execution** by the executor and threaded into the
  ``AdapterBundle``; it is never a process global.
- The executor's cancellation watcher calls ``cancel()`` when it observes
  ``cancel_requested`` in the DB; the running ``asyncio.Task`` sees
  ``CancelledError`` propagate up.
"""

from __future__ import annotations

import asyncio


class CancellationToken:
    """A lightweight cooperative cancel flag.

    Thread-safe in the asyncio sense (single event loop): ``cancel()`` and
    ``is_cancelled`` only touch the flag / an asyncio event.
    """

    __slots__ = ("_cancelled", "_event", "_reason")

    def __init__(self) -> None:
        self._cancelled: bool = False
        self._reason: str = ""
        # Used by ``wait()``; lazily created on first await so we don't need a
        # running loop at construction time (tokens are often built in sync code).
        self._event: asyncio.Event | None = None

    # ── state queries ───────────────────────────────────────────────────────

    @property
    def is_cancelled(self) -> bool:
        return self._cancelled

    @property
    def reason(self) -> str:
        return self._reason

    # ── transitions ─────────────────────────────────────────────────────────

    def cancel(self, reason: str = "cancelled") -> None:
        """Flip the flag and wake any coroutine blocked on ``wait()``."""
        self._cancelled = True
        self._reason = reason
        if self._event is not None:
            self._event.set()

    # ── checked-call helpers ────────────────────────────────────────────────

    def raise_if_cancelled(self) -> None:
        """Raise :class:`ResearchCancelled` if the token has been cancelled.

        Sprinkle at the top of every loop iteration in vendored / runner code:
        search fan-out, context compression, report chunking, graph nodes.
        """
        if self._cancelled:
            raise ResearchCancelled(self._reason)

    async def wait(self, timeout: float | None = None) -> None:
        """Block until cancelled (or ``timeout`` elapses).

        Useful when a node needs to *park* until either cancellation or an
        external resume signal — though for human-review interrupts the
        LangGraph ``interrupt()`` mechanism is preferred.
        """
        if self._cancelled:
            return
        if self._event is None:
            self._event = asyncio.Event()
        try:
            await asyncio.wait_for(self._event.wait(), timeout=timeout)
        except TimeoutError:
            pass

    # ── context-manager sugar (check on enter/exit of a scoped unit) ────────

    def __enter__(self) -> CancellationToken:
        self.raise_if_cancelled()
        return self

    def __exit__(self, *exc: object) -> None:
        self.raise_if_cancelled()


class ResearchCancelled(Exception):
    """Raised when a :class:`CancellationToken` is flipped mid-run.

    Deliberately **not** a subclass of ``asyncio.CancelledError``: the latter is
    used by the executor to kill the task object; ``ResearchCancelled`` is the
    *checked* variant that lets adapters/runners unwind cleanly and record a
    ``cancelled`` terminal state rather than an error.
    """

    def __init__(self, reason: str = "cancelled") -> None:
        super().__init__(reason)
        self.reason = reason


class ResearchAwaitingInput(Exception):
    """Signal a durable human-review pause to the job executor.

    This is intentionally separate from cancellation: the runner has completed
    a safe checkpoint and the executor must preserve its plan, release the
    worker lease, and let a later ``resume`` job continue from that checkpoint.
    """

    def __init__(self, interrupt_id: str, plan: dict) -> None:
        super().__init__(f"Research is awaiting input at {interrupt_id}")
        self.interrupt_id = interrupt_id
        self.plan = plan


def raise_if_cancelled_or_return(value: CancellationToken | None) -> None:
    """Convenience for adapters that receive an optional token."""
    if value is not None:
        value.raise_if_cancelled()


__all__ = [
    "CancellationToken",
    "ResearchAwaitingInput",
    "ResearchCancelled",
    "raise_if_cancelled_or_return",
]