import asyncio
import contextvars
from functools import wraps
import logging
from typing import ParamSpec, TypeVar, Any
from collections.abc import Callable, Coroutine
from celery import signals
from acelery.runner import Runner

logger = logging.getLogger(__name__)

EventLoopFactory = Callable[[], asyncio.AbstractEventLoop]

P = ParamSpec("P")
Ret = TypeVar("Ret")

_runner: Runner | None = None


@signals.worker_process_init.connect
def setup_worker_process(**kwargs: Any) -> None:  # noqa ARG001
    global _runner
    _runner = Runner()
    _runner.__enter__()


@signals.worker_process_shutdown.connect
def cleanup_worker_process(**kwargs: Any) -> None:  # noqa ARG001
    global _runner
    if _runner is not None:
        _runner.close()
        _runner = None


async def cancel_leftover_tasks():
    running_tasks = asyncio.all_tasks()
    running_tasks.remove(asyncio.current_task())
    if running_tasks:
        tasks_report = "\n".join(str(task) for task in running_tasks)
        logger.error(
            "Detected %s running coroutines after task completion: \n%s",
            len(running_tasks),
            tasks_report,
        )
        for task in running_tasks:
            logger.warning("Cancelling task: %s", task)
            task.cancel()


def async_task(
    *, dedicated_runner: bool = False
) -> Callable[[Callable[P, Coroutine[Any, Any, Ret]]], Callable[P, Ret]]:
    """
    Decorator to run a coroutine function as a task
    The coroutine will prefer to execute inside the worker's global asyncio runner, or create a dedicated runner if needed.
    :param dedicated_runner: Use a dedicated asyncio runner for the task instead of the worker's global runner (for isolation, extra ±10µs)
    :return: Decorated task function
    """

    def decorator(func: Callable[P, Coroutine[Any, Any, Ret]]) -> Callable[P, Ret]:
        if not asyncio.iscoroutinefunction(func):
            raise ValueError(
                "Task %s is decorated with @async_task but is not a coroutine function",
                func.__qualname__,
            )

        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> Ret:
            async def _run() -> Ret:
                result = await func(*args, **kwargs)
                await cancel_leftover_tasks()
                return result

            ctx = contextvars.copy_context()

            global _runner
            if not dedicated_runner and _runner is not None:
                return _runner.run(_run(), context=ctx)
            else:
                logger.info(
                    "Running task %s with a dedicated asyncio runner", func.__qualname__
                )
                return asyncio.run(_run())

        return wrapper

    return decorator
