import asyncio
from asyncio import CancelledError

from chancy import Reference
from chancy.executors.base import Executor
from chancy.job import QueuedJob, Limit


class AsyncExecutor(Executor):
    """
    An Executor which uses asyncio to run its jobs in the main event loop.

    This executor is useful for running large numbers of IO-bound jobs, as it
    can run many jobs concurrently without blocking the main event loop and
    without the high overhead of new processes or threads. However, it is not
    suitable for CPU-bound jobs, as it will block the main event loop and
    prevent other jobs & queues from running.

    To use this executor, simply pass the import path to this class in the
    ``executor`` field of your queue configuration or use the
    :class:`~chancy.app.Chancy.Executor` shortcut:

    .. code-block:: python

        async with Chancy(dsn="postgresql://localhost/postgres") as chancy:
            await chancy.declare(
                Queue(
                    name="default",
                    executor=Chancy.Executor.Async,
                )
            )
    """

    def __init__(self, worker, queue):
        super().__init__(worker, queue)
        self.jobs: dict[asyncio.Task, QueuedJob] = {}

    async def push(self, job: QueuedJob):
        task = asyncio.create_task(self._job_wrapper(job))
        self.jobs[task] = job
        task.add_done_callback(self.jobs.pop)

    def __len__(self):
        return len(self.jobs)

    async def _job_wrapper(self, job: QueuedJob):
        try:
            func, kwargs = Executor.get_function_and_kwargs(job)
            if not asyncio.iscoroutinefunction(func):
                raise ValueError(
                    f"Function {job.func!r} is not an async function, which is"
                    f" required for the AsyncExecutor. Please use the"
                    f" ThreadedExecutor or ProcessExecutor instead."
                )

            timeout = next(
                (
                    limit.value
                    for limit in job.limits
                    if limit.type_ == Limit.Type.TIME
                ),
                None,
            )

            async with asyncio.timeout(timeout):
                await func(**kwargs)
            await self.job_completed(job)
        except (Exception, CancelledError) as exc:
            await self.job_completed(job, exc)

    async def cancel(self, ref: Reference):
        for task, job in self.jobs.items():
            if job.id == ref.identifier:
                task.cancel()
                return

    async def stop(self):
        """
        Stop the executor, giving it a chance to clean up any resources it
        may have allocated to running jobs.
        """
        for task in self.jobs:
            task.cancel()
        await asyncio.gather(*self.jobs)
