from __future__ import annotations

import logging
import typing as t
from abc import ABC
from abc import abstractmethod

import attr
from simple_di import Provide
from simple_di import inject

from ...exceptions import StateException
from ..configuration.containers import BentoMLContainer
from ..models.model import Model
from ..tag import validate_tag_str
from ..utils import deprecated
from ..utils import first_not_none
from .runnable import Runnable
from .runner_handle import DummyRunnerHandle
from .runner_handle import RunnerHandle
from .strategy import DefaultStrategy
from .strategy import Strategy

if t.TYPE_CHECKING:
    from ...triton import Runner as TritonRunner
    from .runnable import RunnableMethodConfig

    # only use ParamSpec in type checking, as it's only in 3.10
    P = t.ParamSpec("P")
    ListModel = list[Model]
else:
    P = t.TypeVar("P")
    ListModel = list

T = t.TypeVar("T", bound=Runnable)
R = t.TypeVar("R")


logger = logging.getLogger(__name__)

object_setattr = object.__setattr__


@attr.frozen(slots=False)
class RunnerMethod(t.Generic[T, P, R]):
    runner: Runner | TritonRunner
    name: str
    config: RunnableMethodConfig
    max_batch_size: int
    max_latency_ms: int
    doc: str | None = None

    def run(self, *args: P.args, **kwargs: P.kwargs) -> R:
        return self.runner._runner_handle.run_method(self, *args, **kwargs)

    async def async_run(self, *args: P.args, **kwargs: P.kwargs) -> R:
        return await self.runner._runner_handle.async_run_method(self, *args, **kwargs)

    def async_stream(
        self, *args: P.args, **kwargs: P.kwargs
    ) -> t.AsyncGenerator[str, None]:
        return self.runner._runner_handle.async_stream_method(self, *args, **kwargs)


def _to_lower_name(name: str) -> str:
    lname = name.lower()
    if name != lname:
        logger.warning("Converting runner name '%s' to lowercase: '%s'", name, lname)

    return lname


def _validate_name(_: t.Any, attr: attr.Attribute[str], value: str):
    try:
        validate_tag_str(value)
    except ValueError as e:
        # TODO: link to tag validation documentation
        raise ValueError(
            f"Runner name '{value}' is not valid; it must be a valid BentoML Tag name."
        ) from e


@attr.define(slots=False, frozen=True)
class AbstractRunner(ABC):
    name: str = attr.field(converter=_to_lower_name, validator=_validate_name)
    models: list[Model] = attr.field(
        converter=attr.converters.default_if_none(factory=list),
        validator=attr.validators.deep_iterable(
            attr.validators.instance_of(Model),
            iterable_validator=attr.validators.instance_of(ListModel),
        ),
    )
    resource_config: dict[str, t.Any]
    runnable_class: type[Runnable]
    embedded: bool

    @abstractmethod
    def init_local(self, quiet: bool = False) -> None:
        """
        Initialize local runnable instance, for testing and debugging only.

        Args:
            quiet: if True, no logs will be printed
        """

    @abstractmethod
    def init_client(
        self,
        handle_class: type[RunnerHandle] | None = None,
        *args: t.Any,
        **kwargs: t.Any,
    ):
        """
        Initialize client for a remote runner instance. To be used within API server instance.
        """


@deprecated(suggestion="Please upgrade to new style services.")
@attr.define(slots=False, frozen=True, eq=False, init=False)
class Runner(AbstractRunner):
    if t.TYPE_CHECKING:
        # This will be set by __init__. This is for type checking only.
        run: t.Callable[..., t.Any]
        async_run: t.Callable[..., t.Awaitable[t.Any]]

        # the following annotations hacks around the fact that Runner does not
        # have information about signatures at runtime.
        @t.overload
        def __getattr__(
            self, item: t.Literal["__attrs_init__"]
        ) -> t.Callable[..., None]:  # type: ignore
            ...

        @t.overload
        def __getattr__(
            self, item: t.LiteralString
        ) -> RunnerMethod[t.Any, P, t.Any]: ...

        def __getattr__(self, item: str) -> t.Any: ...

    runner_methods: list[RunnerMethod[t.Any, t.Any, t.Any]]
    scheduling_strategy: type[Strategy]
    workers_per_resource: int | float = 1
    runnable_init_params: dict[str, t.Any] = attr.field(
        default=None, converter=attr.converters.default_if_none(factory=dict)
    )
    _runner_handle: RunnerHandle = attr.field(init=False, factory=DummyRunnerHandle)

    def _set_handle(
        self, handle_class: type[RunnerHandle], *args: t.Any, **kwargs: t.Any
    ) -> None:
        if not isinstance(self._runner_handle, DummyRunnerHandle):
            raise StateException("Runner already initialized")

        runner_handle = handle_class(self, *args, **kwargs)
        object_setattr(self, "_runner_handle", runner_handle)

    @inject
    async def runner_handle_is_ready(
        self,
        timeout: int = Provide[BentoMLContainer.api_server_config.runner_probe.timeout],
    ) -> bool:
        """
        Check if given runner handle is ready. This will be used as readiness probe in Kubernetes.
        """
        return await self._runner_handle.is_ready(timeout)

    def __init__(
        self,
        runnable_class: type[Runnable],
        *,
        runnable_init_params: dict[str, t.Any] | None = None,
        name: str | None = None,
        scheduling_strategy: type[Strategy] = DefaultStrategy,
        models: list[Model] | None = None,
        max_batch_size: int | None = None,
        max_latency_ms: int | None = None,
        method_configs: dict[str, dict[str, int]] | None = None,
        embedded: bool = False,
    ) -> None:
        """

        Runner represents a unit of computation that can be executed on a remote Python worker and scales independently
        See https://docs.bentoml.com/en/latest/concepts/runner.html for more details.

        Args:
            runnable_class: Runnable class that can be executed on a remote Python worker.
            runnable_init_params: Parameters to be passed to the runnable class constructor ``__init__``.
            name: Given a name for this runner. If not provided, name will be generated from the runnable class name.
                  Note that all name will be converted to lowercase and validate to be a valid BentoML Tag name.
            scheduling_strategy: A strategy class that implements the scheduling logic for this runner. If not provided,
                                 use the default strategy. Strategy will respect ``Runnable.SUPPORTED_RESOURCES`` as well as
                                 ``Runnable.SUPPORTS_CPU_MULTI_THREADING``.
            models: An optional list composed of ``bentoml.Model`` instances.
            max_batch_size: Max batch size config for dynamic batching. If not provided, use the default value from
                            configuration.
            max_latency_ms: Max latency config for dynamic batching. If not provided, use the default value from
                            configuration.
            method_configs: A dictionary per method config for this given Runner signatures.

        Returns:
            :obj:`bentoml.Runner`: A Runner instance.
        """

        if name is None:
            name = runnable_class.__name__.lower()
            logger.warning(
                "Using lowercased runnable class name '%s' for runner.", name
            )

        runners_config = BentoMLContainer.config.runners.get()
        # If given runner is configured, then use it. Otherwise use the default configuration.
        if name in runners_config:
            config = runners_config[name]
        else:
            config = runners_config

        if models is None:
            models = []

        runner_method_map: dict[str, RunnerMethod[t.Any, t.Any, t.Any]] = {}
        method_configs = {} if method_configs is None else method_configs

        if runnable_class.bentoml_runnable_methods__ is None:
            raise ValueError(
                f"Runnable class '{runnable_class.__name__}' has no methods!"
            )

        for method_name, method in runnable_class.bentoml_runnable_methods__.items():
            if not config["batching"]["enabled"]:
                method.config.batchable = False

            method_max_batch_size = None
            method_max_latency_ms = None
            if method_name in method_configs:
                method_max_batch_size = method_configs[method_name].get(
                    "max_batch_size"
                )
                method_max_latency_ms = method_configs[method_name].get(
                    "max_latency_ms"
                )

            runner_method_map[method_name] = RunnerMethod(
                runner=self,
                name=method_name,
                config=method.config,
                doc=method.func.__doc__,
                max_batch_size=first_not_none(
                    method_max_batch_size,
                    max_batch_size,
                    default=config["batching"]["max_batch_size"],
                ),
                max_latency_ms=first_not_none(
                    method_max_latency_ms,
                    max_latency_ms,
                    default=config["batching"]["max_latency_ms"],
                ),
            )

        self.__attrs_init__(
            name=name,
            models=models,
            runnable_class=runnable_class,
            runnable_init_params=runnable_init_params,
            resource_config=config["resources"],
            workers_per_resource=config.get("workers_per_resource", 1),
            runner_methods=list(runner_method_map.values()),
            scheduling_strategy=scheduling_strategy,
            embedded=embedded,
        )

        # Choose the default method:
        #  1. if there's only one method, it will be set as default
        #  2. if there's a method named "__call__", it will be set as default
        #  3. otherwise, there's no default method
        if len(runner_method_map) == 1:
            default_method = next(iter(runner_method_map.values()))
            logger.debug(
                "Default runner method set to '%s', it can be accessed both via 'runner.run' and 'runner.%s.async_run'.",
                default_method.name,
                default_method.name,
            )
        elif "__call__" in runner_method_map:
            default_method = runner_method_map["__call__"]
            logger.debug(
                "Default runner method set to '__call__', it can be accessed via 'runner.run' or 'runner.async_run'."
            )
        else:
            default_method = None
            logger.debug(
                "No default method found for Runner '%s', all method access needs to be in the form of 'runner.{method}.run'.",
                name,
            )

        # set default run method entrypoint
        if default_method is not None:
            object.__setattr__(self, "run", default_method.run)
            object.__setattr__(self, "async_run", default_method.async_run)

        # set all run method entrypoint
        for runner_method in self.runner_methods:
            object.__setattr__(self, runner_method.name, runner_method)

    def init_local(self, quiet: bool = False) -> None:
        if not quiet:
            logger.warning(
                "'Runner.init_local' is for debugging and testing only. Make sure to remove it before deploying to production."
            )

        from .runner_handle.local import LocalRunnerRef

        try:
            self._set_handle(LocalRunnerRef)
        except Exception as e:
            import traceback

            logger.error(
                "An exception occurred while instantiating runner '%s', see details below:",
                self.name,
            )
            logger.error(traceback.format_exc())

            raise e

    def init_client(
        self,
        handle_class: type[RunnerHandle] | None = None,
        *args: t.Any,
        **kwargs: t.Any,
    ):
        if handle_class is None:
            from .runner_handle.remote import RemoteRunnerClient

            self._set_handle(RemoteRunnerClient)
        else:
            self._set_handle(handle_class, *args, **kwargs)

    def destroy(self):
        """
        Destroy the runner. This is called when the runner is no longer needed.
        Currently used under ``on_shutdown`` event of the BentoML server.
        """
        object_setattr(self, "_runner_handle", DummyRunnerHandle())

    @property
    def scheduled_worker_count(self) -> int:
        return self.scheduling_strategy.get_worker_count(
            self.runnable_class,
            self.resource_config,
            self.workers_per_resource,
        )

    @property
    def scheduled_worker_env_map(self) -> dict[int, dict[str, t.Any]]:
        return {
            worker_id: self.scheduling_strategy.get_worker_env(
                self.runnable_class,
                self.resource_config,
                self.workers_per_resource,
                worker_id,
            )
            for worker_id in range(self.scheduled_worker_count)
        }
