# 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.
# pylint: disable=invalid-name, missing-function-docstring
"""A utility method to run a TVM module on a remote device."""

from collections.abc import Callable
from typing import TYPE_CHECKING, Literal, Optional, Union

if TYPE_CHECKING:
    import numpy as np

    from tvm.runtime import Device, Module, Tensor
    from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig

# pylint: disable=import-outside-toplevel,protected-access


def _args_to_device(args, device):
    import numpy as np

    from tvm.runtime import Tensor, empty

    uploaded_args = []
    for arg in args:
        if isinstance(arg, np.ndarray | Tensor):
            uploaded_args.append(empty(arg.shape, dtype=arg.dtype, device=device).copyfrom(arg))
        elif isinstance(arg, int | float):
            uploaded_args.append(arg)
        else:
            raise ValueError(f"Unsupported input type: {type(arg)}")
    return uploaded_args


def _args_to_numpy(args):
    from tvm.runtime import Tensor

    downloaded_args = []
    for arg in args:
        if isinstance(arg, Tensor):
            downloaded_args.append(arg.numpy())
        else:
            downloaded_args.append(arg)
    return downloaded_args


def _normalize_export_func(export_func, output_format) -> tuple[Callable, str]:
    from tvm.support import ndk, tar

    def export_with(func):
        return lambda mod, path: mod.export_library(path, fcompile=func)

    if export_func == "tar":
        export_func = export_with(tar.tar)
        output_format = "tar"
    elif export_func == "ndk":
        export_func = export_with(ndk.create_shared)
        output_format = "so"
    elif callable(export_func):
        if output_format is None:
            raise ValueError("output_format must be specified if `export_func` is callable")
    else:
        raise ValueError(f"Unsupported export_func: {export_func}")
    return export_func, output_format


def local_run(  # pylint: disable=too-many-arguments,too-many-locals
    mod: "Module",
    device_type: str,
    args: list[Union["np.ndarray", "Tensor", int, float]],
    evaluator_config: Optional["EvaluatorConfig"] = None,
    export_func: Callable[["Module", str], None] | Literal["tar", "ndk"] = "tar",
    output_format: str | None = None,
):
    """Run a TVM module on a local device.

    Parameters
    ----------
    mod : Module
        The TVM module to run.
    device_type : str
        The device type to run the module on.
    args : List[Union[np.ndarray, Tensor, int, float]]
        The arguments to be fed to the module.
    evaluator_config : Optional[EvaluatorConfig]
        The evaluator configuration to use.
    export_func : Union[Callable[Module, str], Literal["tar", "ndk"]]
        The function to export the module to a file.
        If callable, it must be a function that takes two arguments: the module to export and the
        path to export to.
        If "tar", the module will be exported to a tar file.
        If "ndk", the module will be exported to a shared library.
    output_format : Optional[str]
        The format of the exported module.
        If not specified, it will be inferred from the `export_func` argument.

    Returns
    -------
    args : List[Union[np.ndarray, Tensor, int, float]]
        The results of running the module.
    profile_result : tvm.runtime.BenchmarkResult
        The profiling result of running the module.
    """
    import os.path as osp
    import tempfile

    from tvm.runtime import device, load_module
    from tvm.s_tir.meta_schedule.runner import EvaluatorConfig

    evaluator_config = EvaluatorConfig._normalized(evaluator_config)
    export_func, output_format = _normalize_export_func(export_func, output_format)

    with tempfile.TemporaryDirectory() as tmp_dir:
        artifact_path = osp.join(tmp_dir, "tvm_tmp_mod." + output_format)
        export_func(mod, artifact_path)
        device: Device = device(device_type, 0)

        args = _args_to_device(args, device)
        remote_mod = load_module(artifact_path)
        profile_result = remote_mod.time_evaluator(
            func_name=remote_mod.entry_name,
            dev=device,
            number=evaluator_config.number,
            repeat=evaluator_config.repeat,
            min_repeat_ms=evaluator_config.min_repeat_ms,
            f_preproc="cache_flush_cpu_non_first_arg"
            if evaluator_config.enable_cpu_cache_flush
            else "",
        )(*args)
        remote_mod(*args)
        args = _args_to_numpy(args)

    return args, profile_result


def rpc_run(  # pylint: disable=too-many-arguments,too-many-locals
    mod: "Module",
    device_type: str,
    args: list[Union["np.ndarray", "Tensor", int, float]],
    evaluator_config: Optional["EvaluatorConfig"] = None,
    rpc_config: Optional["RPCConfig"] = None,
    export_func: Callable[["Module", str], None] | Literal["tar", "ndk"] = "tar",
    output_format: str | None = None,
):
    """Run a TVM module on a remote device.

    Parameters
    ----------
    mod : Module
        The TVM module to run.
    device_type : str
        The device type to run the module on.
    args : List[Union[np.ndarray, Tensor, int, float]]
        The arguments to be fed to the module.
    evaluator_config : Optional[EvaluatorConfig]
        The evaluator configuration to use.
    rpc_config : Optional[RPCConfig]
        The RPC configuration to connect to the remote device.
        If not specified, the default RPC configuration will be used, which reads the following
        environment variables:
        - TVM_TRACKER_HOST
        - TVM_TRACKER_PORT
        - TVM_TRACKER_KEY
    export_func : Union[Callable[Module, str], Literal["tar", "ndk"]]
        The function to export the module to a file.
        If callable, it must be a function that takes two arguments: the module to export and the
        path to export to.
        If "tar", the module will be exported to a tar file.
        If "ndk", the module will be exported to a shared library.
    output_format : Optional[str]
        The format of the exported module.
        If not specified, it will be inferred from the `export_func` argument.

    Returns
    -------
    args : List[Union[np.ndarray, Tensor, int, float]]
        The results of running the module.
    profile_result : tvm.runtime.BenchmarkResult
        The profiling result of running the module.
    """

    import os.path as osp
    import tempfile

    from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig

    evaluator_config = EvaluatorConfig._normalized(evaluator_config)
    rpc_config = RPCConfig._normalized(rpc_config)
    export_func, output_format = _normalize_export_func(export_func, output_format)

    with tempfile.TemporaryDirectory() as tmp_dir:
        artifact_path = osp.join(tmp_dir, "tvm_tmp_mod." + output_format)
        _, remote_path = osp.split(artifact_path)
        session = rpc_config.connect_server()
        device: Device = session.device(device_type, 0)

        export_func(mod, artifact_path)
        try:
            session.upload(artifact_path, remote_path)
            args = _args_to_device(args, device)
            remote_mod = session.load_module(remote_path)
            profile_result = remote_mod.time_evaluator(
                func_name=remote_mod.entry_name,
                dev=device,
                number=evaluator_config.number,
                repeat=evaluator_config.repeat,
                min_repeat_ms=evaluator_config.min_repeat_ms,
                f_preproc="cache_flush_cpu_non_first_arg"
                if evaluator_config.enable_cpu_cache_flush
                else "",
            )(*args)
            remote_mod(*args)
            args = _args_to_numpy(args)
        finally:
            session.remove(remote_path)
            session.remove(remote_path + "." + output_format)
            session.remove("")

    return args, profile_result