# 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.
# ruff: noqa: F401

# pylint: disable=invalid-name, unused-import, import-outside-toplevel, inconsistent-return-statements
"""Runtime Module namespace."""

import json
import os
import struct
from collections.abc import Sequence

import numpy as np
from tvm_ffi import (
    Module as _Module,
)
from tvm_ffi import (
    load_module as _load_module,
)
from tvm_ffi import (
    register_object as _register_object,
)
from tvm_ffi import (
    system_lib,
)

from tvm.base import _RUNTIME_ONLY
from tvm.libinfo import find_include_path

from . import _ffi_api


class BenchmarkResult:
    """Runtimes from benchmarking"""

    def __init__(self, results: Sequence[float]):
        """Construct a new BenchmarkResult from a sequence of runtimes.

        Parameters
        ----------
        results : Sequence[float]
            Raw times from benchmarking

        Attributes
        ----------
        min : float
            Minimum runtime in seconds of all results.
        mean : float
            Mean runtime in seconds of all results. If py:meth:`Module.time_evaluator` or
            `benchmark` is called with `number` > 0, then each result is already the mean of a
            `number` of runtimes, so this becomes the mean of means.
        median : float
            Median runtime in seconds of all results. If py:meth:`Module.time_evaluator` is called
            with `number` > 0, then each result is already the mean of a `number` of runtimes, so
            this becomes the median of means.
        max : float
            Maximum runtime in seconds of all results. If py:meth:`Module.time_evaluator` is called
            with `number` > 0, then each result is already the mean of a `number` of runtimes, so
            this becomes the maximum of those means.
        std : float
            Standard deviation in seconds of runtimes. If py:meth:`Module.time_evaluator` is called
            with `number` > 0, then each result is already the mean of a `number` of runtimes, so
            this becomes the standard deviation of means.
        results : Sequence[float]
            The collected runtimes (in seconds). This may be a series of mean runtimes if
            py:meth:`Module.time_evaluator` or `benchmark` was run with `number` > 1.
        """
        self.results = results
        self.mean = np.mean(self.results)
        self.std = np.std(self.results)
        self.median = np.median(self.results)
        self.min = np.min(self.results)
        self.max = np.max(self.results)

    def __repr__(self):
        return (
            f"BenchmarkResult(min={self.min}, mean={self.mean}, median={self.median}, "
            f"max={self.max}, std={self.std}, results={self.results})"
        )

    def __str__(self):
        return (
            f"Execution time summary:\n"
            f"{'mean (ms)':^12} {'median (ms)':^12} {'max (ms)':^12} "
            f"{'min (ms)':^12} {'std (ms)':^12}\n"
            f"{self.mean * 1000:^12.4f} {self.median * 1000:^12.4f} {self.max * 1000:^12.4f} "
            f"{self.min * 1000:^12.4f} {self.std * 1000:^12.4f}"
            "               "
        )


# override the Module class in ffi.Module
@_register_object("ffi.Module")
class Module(_Module):
    """Runtime Module."""

    def _collect_from_import_tree(self, filter_func):
        """Helper function to collect modules from the tree matching a filter_func, then return it.

        Parameters
        ----------
        filter_func : Callable[[Module], bool]
            A function which is invoked for each Module discovered in the import tree (including
            self).

        Returns
        -------
        list[Module] :
            A list of matching Module.
        """
        visited, stack, dso_modules = set(), [], []
        # append root module
        visited.add(self)
        stack.append(self)
        while stack:
            module = stack.pop()
            assert module.is_compilation_exportable() or module.is_binary_serializable(), (
                f"Module {module.kind} should be either dso exportable or binary serializable."
            )

            if filter_func(module):
                dso_modules.append(module)
            for m in module.imports:
                if m not in visited:
                    visited.add(m)
                    stack.append(m)
        return dso_modules

    def _collect_dso_modules(self):
        """Collect all compilation exportable modules from the import tree."""
        return self._collect_from_import_tree(lambda m: m.is_compilation_exportable())

    def export_library(
        self,
        file_name,
        *,
        fcompile=None,
        fpack_imports=None,
        addons=None,
        workspace_dir=None,
        **kwargs,
    ):
        """
        Export the module and all imported modules into a single device library.

        This function only works on host LLVM modules, other runtime::Module
        subclasses will work with this API but they must support implement
        the save and load mechanisms of modules completely including saving
        from streams and files. This will pack your non-shared library module
        into a single shared library which can later be loaded by TVM.

        Parameters
        ----------
        file_name : str
            The name of the shared library.

        fcompile : function(target, file_list, kwargs), optional
            The compilation function to use create the final library object during
            export.

            For example, when fcompile=_cc.create_shared, or when it is not supplied but
            module is "llvm," this is used to link all produced artifacts
            into a final dynamic library.

            This behavior is controlled by the type of object exported.
            If fcompile has attribute object_format, will compile host library
            to that format. Otherwise, will use default format "o".

        fpack_imports: function(mod: runtime.Module, is_system_lib: bool, symbol_prefix: str,
                                workspace_dir: str) -> str
            Function used to pack imported modules from `mod` into a file suitable for passing
            to fcompile as an input file. The result can be a C source, or an .o object file,
            or any other file that the fcompile function can handle. The function returns the
            name of the created file.

            If not provided, the imported modules will be serialized either via packing to an
            LLVM module, or to a C source file.

        workspace_dir : str, optional
            The path of the directory used to create the intermediate
            artifacts when exporting the module.
            If this is not provided a temporary dir will be created.

        kwargs : dict, optional
            Additional arguments passed to fcompile

        Returns
        -------
        result of fcompile()  : unknown, optional
            If the compilation function returns an artifact it would be returned via
            export_library, if any.
        """
        # NOTE: this function depends on contrib library features
        # which are only available in when TVM function is available.
        if _RUNTIME_ONLY:
            raise RuntimeError("Cannot call export_library in runtime only mode")

        # Extra dependencies during runtime.
        from pathlib import Path

        from tvm.contrib import cc as _cc
        from tvm.contrib import tar as _tar
        from tvm.contrib import tvmjs as _tvmjs
        from tvm.contrib import utils as _utils

        if isinstance(file_name, Path):
            file_name = str(file_name)

        modules = self._collect_dso_modules()
        if workspace_dir is None:
            temp = _utils.tempdir()
            workspace_dir = temp.temp_dir
        files = addons if addons else []
        is_system_lib = False
        has_c_module = False
        system_lib_prefix = None
        llvm_target = None
        global_object_format = "o"

        def get_source_format_from_module(module):
            for fmt in module.get_write_formats():
                if fmt in ["c", "cc", "cpp", "cu"]:
                    return fmt
            raise ValueError(f"Module {module.kind} does not exporting to c, cc, cpp or cu.")

        for index, module in enumerate(modules):
            if fcompile is not None and hasattr(fcompile, "object_format"):
                if module.kind == "c":
                    object_format = get_source_format_from_module(module)
                    has_c_module = True
                else:
                    global_object_format = object_format = fcompile.object_format
            else:
                if module.kind == "c":
                    if len(module.get_write_formats()) > 0:
                        object_format = get_source_format_from_module(module)
                    else:
                        object_format = "c"
                    if "cc" in kwargs:
                        if kwargs["cc"] == "nvcc":
                            object_format = "cu"
                    has_c_module = True
                else:
                    assert module.is_compilation_exportable()
                    global_object_format = object_format = "o"

            path_obj = os.path.join(workspace_dir, f"lib{index}.{object_format}")
            module.write_to_file(path_obj)
            files.append(path_obj)
            if module.kind == "llvm":
                is_system_lib = module.get_function("__tvm_is_system_module")()
                llvm_target = module.get_function("_get_target_string")()
                system_lib_prefix = module.get_function("__tvm_get_system_lib_prefix")()

        if not fcompile:
            if file_name.endswith(".tar"):
                fcompile = _tar.tar
            elif file_name.endswith(".wasm"):
                fcompile = _tvmjs.create_tvmjs_wasm
            else:
                fcompile = _cc.create_shared

        if llvm_target is None and hasattr(fcompile, "get_target_triple"):
            triple = fcompile.get_target_triple()
            assert triple, "Target triple should not be empty"
            llvm_target = json.dumps({"kind": "llvm", "mtriple": triple.strip()})

        if getattr(fcompile, "need_system_lib", False) and not is_system_lib:
            raise ValueError(f"{fcompile!s} need --system-lib option")

        if self.imports:
            pack_lib_prefix = system_lib_prefix if system_lib_prefix else ""

            if fpack_imports is not None:
                path_out = fpack_imports(self, is_system_lib, pack_lib_prefix, workspace_dir)
                files.append(path_out)
            elif _ffi_api.RuntimeEnabled("llvm") and llvm_target:
                path_obj = os.path.join(
                    workspace_dir, f"{pack_lib_prefix}devc.{global_object_format}"
                )
                m = _ffi_api.ModulePackImportsToLLVM(
                    self, is_system_lib, llvm_target, pack_lib_prefix
                )
                m.write_to_file(path_obj)
                files.append(path_obj)
            else:
                path_cc = os.path.join(workspace_dir, f"{pack_lib_prefix}devc.c")
                with open(path_cc, "w") as f:
                    f.write(_ffi_api.ModulePackImportsToC(self, is_system_lib, pack_lib_prefix))
                files.append(path_cc)

        # The imports could contain a c module but the object format could be tar
        # Thus, it would not recognize the following include paths as options
        # which are there assuming a c compiler is the fcompile.
        if has_c_module and not file_name.endswith(".tar"):
            options = []
            if "options" in kwargs:
                opts = kwargs["options"]
                options = opts if isinstance(opts, list | tuple) else [opts]
            opts = options + ["-I" + path for path in find_include_path()]
            kwargs.update({"options": opts})

        return fcompile(file_name, files, **kwargs)

    def time_evaluator(
        self,
        func_name,
        dev,
        number=10,
        repeat=1,
        min_repeat_ms=0,
        limit_zero_time_iterations=100,
        cooldown_interval_ms=0,
        repeats_to_cooldown=1,
        cache_flush_bytes=0,
        f_preproc="",
    ):
        """Get an evaluator that measures time cost of running function.

        Parameters
        ----------
        func_name: str
            The name of the function in the module.

        dev: Device
            The device we should run this function on.

        number: int
            The number of times to run this function for taking average.
            We call these runs as one `repeat` of measurement.

        repeat: int, optional
            The number of times to repeat the measurement.
            In total, the function will be invoked (1 + number x repeat) times,
            where the first one is warm up and will be discarded.
            The returned result contains `repeat` costs,
            each of which is an average of `number` costs.

        min_repeat_ms: int, optional
            The minimum duration of one `repeat` in milliseconds.
            By default, one `repeat` contains `number` runs. If this parameter is set,
            the parameters `number` will be dynamically adjusted to meet the
            minimum duration requirement of one `repeat`.
            i.e., When the run time of one `repeat` falls below this time, the `number` parameter
            will be automatically increased.

        limit_zero_time_iterations: int, optional
            The maximum number of repeats when measured time is equal to 0.
            It helps to avoid hanging during measurements.

        cooldown_interval_ms: int, optional
            The cooldown interval in milliseconds between the number of repeats defined by
            `repeats_to_cooldown`.

        repeats_to_cooldown: int, optional
            The number of repeats before the cooldown is activated.

        cache_flush_bytes: int, optional
            The number of bytes to flush from the cache before each repeat.

        f_preproc: str, optional
            The preprocess function name we want to execute before executing the time evaluator.

        Note
        ----
        The function will be invoked  (1 + number x repeat) times,
        with the first call discarded in case there is lazy initialization.

        Returns
        -------
        ftimer : function
            The function that takes same argument as func and returns a BenchmarkResult.
            The ProfileResult reports `repeat` time costs in seconds.
        """
        try:
            feval = _ffi_api.RPCTimeEvaluator(
                self,
                func_name,
                dev.dlpack_device_type(),
                dev.index,
                number,
                repeat,
                min_repeat_ms,
                limit_zero_time_iterations,
                cooldown_interval_ms,
                repeats_to_cooldown,
                cache_flush_bytes,
                f_preproc,
            )

            def evaluator(*args):
                """Internal wrapped evaluator."""
                # Wrap feval so we can add more stats in future.
                blob = feval(*args)
                fmt = "@" + ("d" * repeat)
                results = struct.unpack(fmt, blob)
                return BenchmarkResult(results)

            return evaluator
        except NameError:
            raise NameError("time_evaluator is only supported when RPC is enabled")


def load_module(path):
    """Load module from file.

    Parameters
    ----------
    path : str
        The path to the module file.

    Returns
    -------
    module : runtime.Module
        The loaded module

    Note
    ----
    This function will automatically call
    cc.create_shared if the path is in format .o or .tar
    """
    if os.path.isfile(path):
        path = os.path.realpath(path)
    else:
        raise ValueError(f"cannot find file {path}")

    # High level handling for .o and .tar file.
    # We support this to be consistent with RPC module load.
    if path.endswith(".o"):
        # Extra dependencies during runtime.
        from tvm.contrib import cc as _cc

        _cc.create_shared(path + ".so", path)
        path += ".so"
    elif path.endswith(".tar"):
        # Extra dependencies during runtime.
        from tvm.contrib import cc as _cc
        from tvm.contrib import tar as _tar
        from tvm.contrib import utils as _utils

        tar_temp = _utils.tempdir(custom_path=path.replace(".tar", ""))
        _tar.untar(path, tar_temp.temp_dir)
        files = [tar_temp.relpath(x) for x in tar_temp.listdir()]
        _cc.create_shared(path + ".so", files)
        path += ".so"
    # Redirect to the load API
    return _load_module(path)


def load_static_library(path, func_names):
    """Load the .o library at path which implements functions with func_names.
    Unlike the generic load_module the result will remain as a static_library
    and will not be relinked on-the-fly into a .so library."""
    return _ffi_api.ModuleLoadStaticLibrary(path, func_names)


def enabled(target):
    """Whether module runtime is enabled for target

    Parameters
    ----------
    target : str or Dict[str, Any] or tvm.target.Target
        The target device type.

    Returns
    -------
    enabled : bool
        Whether runtime is enabled.

    Examples
    --------
    The following code checks if gpu is enabled.

    >>> tvm.runtime.enabled("gpu")
    """
    if isinstance(target, dict):
        target = target.get("kind", "")
    elif hasattr(target, "kind"):
        target = target.kind.name
    return _ffi_api.RuntimeEnabled(target)


def num_threads() -> int:
    """Get the number of threads in use by the TVM runtime.

    Returns
    -------
    int
        Number of threads in use.
    """
    return _ffi_api.NumThreads()
