# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Local Runner"""

import logging
import subprocess
from collections.abc import Callable
from contextlib import contextmanager

import tvm
from tvm.ir.utils import derived_object
from tvm.support.popen_pool import PopenPoolExecutor

from ....runtime import Device, Module
from ..logging import get_logger
from ..profiler import Profiler
from ..utils import get_global_func_with_default_on_worker
from .config import EvaluatorConfig
from .runner import PyRunner, PyRunnerFuture, RunnerFuture, RunnerInput, RunnerResult
from .utils import (
    T_ARG_INFO_JSON_OBJ_LIST,
    T_ARGUMENT_LIST,
    alloc_argument_common,
    run_evaluator_common,
)

logger = get_logger(__name__)  # pylint: disable=invalid-name


T_ALLOC_ARGUMENT = Callable[  # pylint: disable=invalid-name
    [
        Device,  # The device on the remote
        T_ARG_INFO_JSON_OBJ_LIST,  # The metadata information of the arguments to be allocated
        int,  # The number of repeated allocations to be done
    ],
    list[T_ARGUMENT_LIST],  # A list of argument lists
]
T_RUN_EVALUATOR = Callable[  # pylint: disable=invalid-name
    [
        Module,  # The Module opened on the remote
        Device,  # The device on the remote
        EvaluatorConfig,  # The evaluator configuration
        list[T_ARGUMENT_LIST],  # A list of argument lists
    ],
    list[float],  # A list of running time
]
T_CLEANUP = Callable[  # pylint: disable=invalid-name
    [],
    None,
]


@derived_object
class LocalRunnerFuture(PyRunnerFuture):
    """Local based runner future

    Parameters
    ----------
    res: Optional[List[float]]
        The optional result as a list of float.
    error_message: Optional[str]
        The optional error message.

    Note
    ----
    Only one of the parameters should be None upon the creation
    of LocalRunnerFuture object
    """

    res: list[float] | None
    error_message: str | None

    def __init__(self, res: list[float] | None = None, error_message: str | None = None) -> None:
        """Constructor

        Parameters
        ----------
        res: Optional[List[float]]
            The result of this LocalRunnerFuture
        error_message: Optional[str]
            The stringfied error message of any exception during execution

        """
        super().__init__()
        self.res = res
        self.error_message = error_message

        # sanity check upon the creation of LocalRunnerFuture object
        if (res is None and error_message is None) or (
            res is not None and error_message is not None
        ):
            raise AttributeError(
                "Only one of the two parameters should be None upon the creation"
                "of LocalRunnerFuture object."
            )

    def done(self) -> bool:
        return True

    def result(self) -> RunnerResult:
        return RunnerResult(self.res, self.error_message)


def _worker_func(
    _f_alloc_argument: str | None,
    _f_run_evaluator: str | None,
    _f_cleanup: str | None,
    evaluator_config: EvaluatorConfig,
    alloc_repeat: int,
    artifact_path: str,
    device_type: str,
    args_info: T_ARG_INFO_JSON_OBJ_LIST,
) -> list[float]:
    f_alloc_argument: T_ALLOC_ARGUMENT = get_global_func_with_default_on_worker(
        _f_alloc_argument, default_alloc_argument
    )
    f_run_evaluator: T_RUN_EVALUATOR = get_global_func_with_default_on_worker(
        _f_run_evaluator, default_run_evaluator
    )
    f_cleanup: T_CLEANUP = get_global_func_with_default_on_worker(_f_cleanup, default_cleanup)

    @contextmanager
    def resource_handler():
        try:
            yield
        finally:
            # Final step. Always clean up
            with Profiler.timeit("LocalRunner/cleanup"):
                f_cleanup()

    with resource_handler():
        # Step 1: create the local runtime module
        with Profiler.timeit("LocalRunner/load_module"):
            rt_mod = tvm.runtime.load_module(artifact_path)
        # Step 2: Allocate input arguments
        with Profiler.timeit("LocalRunner/alloc_argument"):
            device = tvm.runtime.device(device_type, 0)
            repeated_args: list[T_ARGUMENT_LIST] = f_alloc_argument(
                device,
                args_info,
                alloc_repeat,
            )
        # Step 3: Run time_evaluator
        with Profiler.timeit("LocalRunner/run_evaluator"):
            costs: list[float] = f_run_evaluator(
                rt_mod,
                device,
                evaluator_config,
                repeated_args,
            )
    return costs


@derived_object
class LocalRunner(PyRunner):
    """Local runner

    Parameters
    ----------
    evaluator_config: EvaluatorConfig
        The evaluator configuration.
    cooldown_sec: float
        The cooldown in seconds.
    alloc_repeat: int
        The number of times to repeat the allocation.
    f_alloc_argument: Optional[str, Callable]
        The function name to allocate the arguments or the function itself.
    f_run_evaluator: Optional[str, Callable]
        The function name to run the evaluator or the function itself.
    f_cleanup: Optional[str, Callable]
        The function name to cleanup the session or the function itself.
    pool: PopenPoolExecutor
        The popen pool executor.

    Attributes
    ----------
    T_ALLOC_ARGUMENT : typing._GenericAlias
        The signature of the function `f_alloc_argument`, which is:

        .. code-block:: python

        def default_alloc_argument(
            device: Device,
            args_info: T_ARG_INFO_JSON_OBJ_LIST,
            alloc_repeat: int,
        ) -> List[T_ARGUMENT_LIST]:
            ...

    T_RUN_EVALUATOR : typing._GenericAlias
        The signature of the function `f_run_evaluator`, which is:

        .. code-block:: python

        def default_run_evaluator(
            rt_mod: Module,
            device: Device,
            evaluator_config: EvaluatorConfig,
            repeated_args: List[T_ARGUMENT_LIST],
        ) -> List[float]:
            ...

    T_CLEANUP : typing._GenericAlias
        The signature of the function `f_cleanup`, which is:

        .. code-block:: python

        def default_cleanup() -> None:
            ...
    """

    timeout_sec: float
    evaluator_config: EvaluatorConfig
    cooldown_sec: float
    alloc_repeat: int

    f_alloc_argument: T_ALLOC_ARGUMENT | str | None
    f_run_evaluator: T_RUN_EVALUATOR | str | None
    f_cleanup: T_CLEANUP | str | None

    pool: PopenPoolExecutor

    def __init__(
        self,
        timeout_sec: float = 30,
        evaluator_config: EvaluatorConfig | None = None,
        cooldown_sec: float = 0.0,
        alloc_repeat: int = 1,
        f_alloc_argument: T_ALLOC_ARGUMENT | str | None = None,
        f_run_evaluator: T_RUN_EVALUATOR | str | None = None,
        f_cleanup: T_CLEANUP | str | None = None,
        initializer: Callable[[], None] | None = None,
    ) -> None:
        """Constructor

        Parameters
        ----------
        timeout_sec: float
            The timeout setting.
        evaluator_config: EvaluatorConfig
            The evaluator configuration.
        cooldown_sec: float
            The cooldown in seconds.
        alloc_repeat: int
            The number of times to random fill the allocation.
        f_alloc_argument: Union[T_ALLOC_ARGUMENT, str, None]
            The function name to allocate the arguments or the function itself.
        f_run_evaluator: Union[T_RUN_EVALUATOR, str, None]
            The function name to run the evaluator or the function itself.
        f_cleanup: Union[T_CLEANUP, str, None]
            The function name to cleanup the session or the function itself.
        initializer: Optional[Callable[[], None]]
            The initializer function.
        """
        super().__init__()
        self.timeout_sec = timeout_sec
        self.evaluator_config = EvaluatorConfig._normalized(evaluator_config)
        self.cooldown_sec = cooldown_sec
        self.alloc_repeat = alloc_repeat
        self.f_alloc_argument = f_alloc_argument
        self.f_run_evaluator = f_run_evaluator
        self.f_cleanup = f_cleanup

        err_path = subprocess.DEVNULL
        if logger.root.level <= logging.DEBUG:
            err_path = subprocess.STDOUT

        logger.info("LocalRunner: max_workers = 1")
        self.pool = PopenPoolExecutor(
            max_workers=1,  # one local worker
            timeout=timeout_sec,
            initializer=initializer,
            stderr=err_path,  # suppress the stderr output
        )
        self._sanity_check()

    def run(self, runner_inputs: list[RunnerInput]) -> list[RunnerFuture]:
        results: list[RunnerFuture] = []
        for runner_input in runner_inputs:
            future = self.pool.submit(
                _worker_func,
                self.f_alloc_argument,
                self.f_run_evaluator,
                self.f_cleanup,
                self.evaluator_config,
                self.alloc_repeat,
                str(runner_input.artifact_path),
                str(runner_input.device_type),
                tuple(arg_info.as_json() for arg_info in runner_input.args_info),
            )
            try:
                result: list[float] = future.result()
                error_message: str = None
            except TimeoutError:
                result = None
                error_message = f"LocalRunner: Timeout, killed after {self.timeout_sec} seconds\n"
            except Exception as exception:  # pylint: disable=broad-except
                result = None
                error_message = "LocalRunner: An exception occurred\n" + str(exception)
            local_future = LocalRunnerFuture(res=result, error_message=error_message)
            results.append(local_future)  # type: ignore
        return results

    def _sanity_check(self) -> None:
        def _check(
            f_alloc_argument,
            f_run_evaluator,
            f_cleanup,
        ) -> None:
            get_global_func_with_default_on_worker(name=f_alloc_argument, default=None)
            get_global_func_with_default_on_worker(name=f_run_evaluator, default=None)
            get_global_func_with_default_on_worker(name=f_cleanup, default=None)

        value = self.pool.submit(
            _check,
            self.f_alloc_argument,
            self.f_run_evaluator,
            self.f_cleanup,
        )
        value.result()


def default_alloc_argument(
    device: Device,
    args_info: T_ARG_INFO_JSON_OBJ_LIST,
    alloc_repeat: int,
) -> list[T_ARGUMENT_LIST]:
    """Default function to allocate the arguments

    Parameters
    ----------
    device: Device
        The device to allocate the arguments
    args_info: T_ARG_INFO_JSON_OBJ_LIST
        The arguments info
    alloc_repeat: int
        The number of times to repeat the allocation

    Returns
    -------
    repeated_args: List[T_ARGUMENT_LIST]
        The allocation args
    """
    f_random_fill = get_global_func_with_default_on_worker(
        name="tvm.contrib.random.random_fill_for_measure", default=None
    )
    return alloc_argument_common(f_random_fill, device, args_info, alloc_repeat)


def default_run_evaluator(
    rt_mod: Module,
    device: Device,
    evaluator_config: EvaluatorConfig,
    repeated_args: list[T_ARGUMENT_LIST],
) -> list[float]:
    """Default function to run the evaluator

    Parameters
    ----------
    rt_mod: Module
        The runtime module
    device: Device
        The device to run the evaluator
    evaluator_config: EvaluatorConfig
        The evaluator config
    repeated_args: List[T_ARGUMENT_LIST]
        The repeated arguments

    Returns
    -------
    costs: List[float]
        The evaluator results
    """
    return run_evaluator_common(rt_mod, device, evaluator_config, repeated_args)


def default_cleanup() -> None:
    """Default function to clean up the session"""
    pass  # pylint: disable=unnecessary-pass


@tvm.register_global_func("s_tir.meta_schedule.runner.get_local_runner")
def get_local_builder() -> LocalRunner:
    """Get the local Runner.

    Returns
    -------
    runner: LocalRunner
        The local runner
    """
    return LocalRunner()