"""Durable in-process task registry.

Both the single-run path (``web.routers.runs``) and the campaign scheduler
(``evaluation.campaign_runner``) drive long-lived ``asyncio`` tasks whose
authority lives in the DB — the in-process handle is only there to cancel or
gracefully stop the task. This module is the one deep seam they share: it owns
the (task, cancel-event) pair per id, wires cleanup, and stops everything on
shutdown. Each caller instantiates its own registry, so id spaces never collide.
"""

import asyncio
import logging
from typing import Awaitable, Callable

_logger = logging.getLogger("agenteval")

# A coroutine factory: given the cancel event the registry owns, return the
# coroutine to run. The event exists before the task starts, so a cancel that
# arrives during startup is never lost.
CoroFactory = Callable[[asyncio.Event], Awaitable[None]]


class TaskRegistry:
    """Tracks live asyncio tasks and their cooperative cancel events by id."""

    def __init__(self) -> None:
        self._tasks: dict[str, asyncio.Task] = {}
        self._cancel_events: dict[str, asyncio.Event] = {}

    def launch(self, id: str, make_coro: CoroFactory) -> asyncio.Task:
        """Start a tracked task for ``id`` (idempotent for a live id).

        Creates the cancel event first, hands it to ``make_coro`` to build the
        coroutine, then schedules the task and wires a done-callback that pops
        both handles when it ends. If ``id`` already has a live task, returns it
        without starting a second coroutine.
        """
        existing = self._tasks.get(id)
        if existing is not None:
            return existing

        cancel = self._cancel_events.setdefault(id, asyncio.Event())
        task = asyncio.create_task(make_coro(cancel), name=id)
        self._tasks[id] = task
        task.add_done_callback(lambda _t, _id=id: self._cleanup(_id))
        return task

    def cancel(self, id: str) -> bool:
        """Signal the cooperative cancel event for ``id``.

        Returns True if a live cancel event was signalled, False if there was
        nothing to signal (e.g. the task already ended or never existed) — the
        caller uses this to decide any fallback such as marking the DB directly.
        """
        event = self._cancel_events.get(id)
        if event is None:
            return False
        event.set()
        return True

    def get(self, id: str) -> "asyncio.Task | None":
        """Return the live task for ``id``, or None if there is none."""
        return self._tasks.get(id)

    def is_running(self, id: str) -> bool:
        return id in self._tasks

    async def shutdown_all(self) -> None:
        """Cooperatively stop, then hard-cancel and await, every live task."""
        for event in list(self._cancel_events.values()):
            event.set()
        tasks = list(self._tasks.values())
        for task in tasks:
            task.cancel()
        for task in tasks:
            # CancelledError is expected (we just cancelled the task) and does
            # not derive from Exception in 3.8+, so list it explicitly.
            try:
                await task
            except (asyncio.CancelledError, Exception) as exc:  # noqa: BLE001
                _logger.debug("task registry shutdown: task ended with %r", exc)
        self._tasks.clear()
        self._cancel_events.clear()

    def _cleanup(self, id: str) -> None:
        self._tasks.pop(id, None)
        self._cancel_events.pop(id, None)