from __future__ import annotations

import functools
import inspect
import warnings
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, Final, Generic, ParamSpec, TypeVar

from typing_extensions import override

from jobify._internal.common.constants import RunMode

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable
    from concurrent.futures import Executor

    from jobify._internal.common.types import LoopFactory
    from jobify._internal.configuration import JobifyConfiguration
    from jobify._internal.inspection import FuncSpec

ReturnT = TypeVar("ReturnT")
ParamsT = ParamSpec("ParamsT")


class RunStrategy(ABC, Generic[ParamsT, ReturnT]):
    __slots__: tuple[str, ...] = ("func",)

    def __init__(self, func: Callable[ParamsT, ReturnT]) -> None:
        self.func: Final = func

    @abstractmethod
    async def __call__(
        self,
        *args: ParamsT.args,
        **kwargs: ParamsT.kwargs,
    ) -> ReturnT:
        raise NotImplementedError


class SyncStrategy(RunStrategy[ParamsT, ReturnT]):
    @override
    async def __call__(
        self,
        *args: ParamsT.args,
        **kwargs: ParamsT.kwargs,
    ) -> ReturnT:
        return self.func(*args, **kwargs)


class AsyncStrategy(RunStrategy[ParamsT, ReturnT]):
    @override
    async def __call__(
        self: AsyncStrategy[ParamsT, Awaitable[ReturnT]],
        *args: ParamsT.args,
        **kwargs: ParamsT.kwargs,
    ) -> ReturnT:
        return await self.func(*args, **kwargs)


class PoolStrategy(RunStrategy[ParamsT, ReturnT]):
    __slots__: tuple[str, ...] = ("executor", "getloop")

    def __init__(
        self,
        func: Callable[ParamsT, ReturnT],
        executor: Executor | None,
        getloop: LoopFactory,
    ) -> None:
        super().__init__(func)
        self.executor: Executor | None = executor
        self.getloop: LoopFactory = getloop

    @override
    async def __call__(
        self,
        *args: ParamsT.args,
        **kwargs: ParamsT.kwargs,
    ) -> ReturnT:
        func_call = functools.partial(self.func, *args, **kwargs)
        return await self.getloop().run_in_executor(self.executor, func_call)


class Runnable(Generic[ReturnT]):
    __slots__: tuple[str, ...] = (
        "bound",
        "func_spec",
        "name",
        "origin_arguments",
        "strategy",
    )

    def __init__(
        self,
        *,
        name: str,
        bound: inspect.BoundArguments,
        strategy: RunStrategy[ParamsT, ReturnT],
        func_spec: FuncSpec[ReturnT],
    ) -> None:
        self.name: str = name
        self.strategy: Final = strategy
        self.bound: inspect.BoundArguments = bound
        self.origin_arguments: dict[str, Any] = bound.arguments.copy()
        self.func_spec: FuncSpec[ReturnT] = func_spec

    def __call__(self) -> Awaitable[ReturnT]:
        return self.strategy(*self.bound.args, **self.bound.kwargs)


def _validate_run_mode(mode: RunMode | None, *, is_async: bool) -> RunMode:
    if is_async:
        if mode in (RunMode.PROCESS, RunMode.THREAD):
            msg = (
                "Async functions are always done in the main loop."
                " This mode (PROCESS/THREAD) is not used."
            )
            warnings.warn(msg, category=RuntimeWarning, stacklevel=3)
        return RunMode.MAIN
    if mode is None:
        return RunMode.THREAD
    return mode


def create_run_strategy(
    func: Callable[ParamsT, ReturnT],
    jobify_config: JobifyConfiguration,
    *,
    mode: RunMode | None,
) -> RunStrategy[ParamsT, ReturnT]:
    # inspect.iscoroutinefunction returns TypeGuard,
    # but we need a regular bool variable
    is_async = bool(inspect.iscoroutinefunction(func))

    mode = _validate_run_mode(mode, is_async=is_async)
    if is_async:
        return AsyncStrategy(func)

    match mode:
        case RunMode.PROCESS:
            processpool = jobify_config.worker_pools.processpool
            return PoolStrategy(func, processpool, jobify_config.getloop)
        case RunMode.THREAD:
            threadpool = jobify_config.worker_pools.threadpool
            return PoolStrategy(func, threadpool, jobify_config.getloop)
        case _:
            return SyncStrategy(func)
