import asyncio
import contextlib
from collections.abc import Coroutine
from typing import Any

from core.utils.coroutines import sentry_wrap


# TODO: this is super basic, only use for really small tasks
class BackgroundTasks:
    def __init__(self):
        self._tasks = set[asyncio.Task[None]]()

    def _safe_pop_task(self, task: asyncio.Task[None]) -> None:
        with contextlib.suppress(KeyError):
            self._tasks.remove(task)

    def add(self, task: Coroutine[Any, Any, None]):
        t = asyncio.create_task(sentry_wrap(task))
        self._tasks.add(t)
        t.add_done_callback(self._safe_pop_task)

    async def wait(self):
        if not self._tasks:
            return
        await asyncio.gather(*self._tasks)
        self._tasks.clear()

    @property
    def task_count(self) -> int:
        return len(self._tasks)


_shared_background_tasks = BackgroundTasks()


def add_background_task(task: Coroutine[Any, Any, None]):
    _shared_background_tasks.add(task)


async def wait_for_background_tasks():
    await _shared_background_tasks.wait()


def active_background_task_count() -> int:
    return _shared_background_tasks.task_count
