from typing import Any, TypeVar
import asyncio
from concurrent.futures import Executor, ThreadPoolExecutor
import functools
import os
from typing import Generic
from collections.abc import Callable
from celery import Signature, Task
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from acelery.result import AIOResult

T = TypeVar("T")


class AIOExecutable(Generic[T]):
    # default is 10 threads per core
    DEFAULT_EXECUTOR = ThreadPoolExecutor(max_workers=10 * (os.cpu_count() or 1))

    def __init__(
        self,
        executable: Signature | Task,
        *,
        return_type: type[T] | None = None,
        executor: Executor = DEFAULT_EXECUTOR,
        loop: asyncio.AbstractEventLoop | None = None,
    ):
        self.executable = executable
        self._executor = executor
        if loop:
            self.__loop = loop
        else:
            try:
                self.__loop = asyncio.get_running_loop()
            except RuntimeError:
                ...  # no running loop, we will lazily get the default loop
        self.return_type = return_type | Any

    @property
    def _loop(self) -> asyncio.AbstractEventLoop:
        if self.__loop is None:
            self.__loop = asyncio.get_event_loop()
        return self.__loop

    async def _run_in_executor(self, fn: Callable, *args, **kwargs):
        if kwargs:
            fn = functools.partial(fn, **kwargs)
        return await self._loop.run_in_executor(self._executor, fn, *args)

    @staticmethod
    def _safe_asyncio_result() -> "type[AIOResult[T]]":
        from acelery.result import AIOResult

        return AIOResult[T]

    async def delay(self, *args, **kwargs) -> "AIOResult[T]":
        celery_result = await self._run_in_executor(
            self.executable.delay,  # type: ignore[attr-defined]
            *args,
            **kwargs,
        )
        return self._safe_asyncio_result()(
            celery_result,
            return_type=self.return_type,
        )

    async def apply_async(self, *args, **kwargs) -> "AIOResult[T]":
        celery_result = await self._run_in_executor(
            self.executable.apply_async,  # type: ignore[attr-defined]
            *args,
            **kwargs,
        )
        return self._safe_asyncio_result()(
            celery_result,
            return_type=self.return_type,
        )

    async def apply(self, *args, **kwargs) -> "AIOResult[T]":
        celery_result = await self._run_in_executor(
            self.executable.apply,  # type: ignore[attr-defined]
            *args,
            **kwargs,
        )
        return self._safe_asyncio_result()(
            celery_result,
            return_type=self.return_type,
        )
