# Copyright 2024 ByteDance and/or its affiliates.
#
# Licensed 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.

from __future__ import annotations

import builtins
import json

import logging
import os
import time
from multiprocessing import Lock
from typing import Any, Callable, Dict, List, Optional

import triton

from protenix.model.tri_attention.autotune_helpers import (
    config_dir,
    config_to_dict,
    device_capability,
    dict_to_config,
)

FILE_LOCK = Lock()


logger = logging.getLogger(__name__)


class Autotuner(triton.runtime.Autotuner):
    """
    Copied from triton -- changed autotune to persist to disk.

    Args:
        fn (Callable): The function to be autotuned.
        arg_names (List[str]): Names of the arguments.
        configs (List[triton.Config]): List of configurations to try.
        key (List[str]): List of argument names to use as cache key.
        reset_to_zero (List[str]): List of argument names to reset to zero.
        restore_value (List[str]): List of argument names to restore.
        pre_hook (Optional[Callable], optional): Pre-hook function. Defaults to None.
        post_hook (Optional[Callable], optional): Post-hook function. Defaults to None.
        prune_configs_by (Optional[Dict], optional): Dictionary of functions to prune configs. Defaults to None.
        warmup (Optional[int], optional): Warmup time in ms. Defaults to None.
        rep (Optional[int], optional): Number of repetitions. Defaults to None.
        use_cuda_graph (bool, optional): Whether to use CUDA graph. Defaults to False.
    """

    def __init__(
        self,
        fn: Callable,
        arg_names: List[str],
        configs: List[triton.Config],
        key: List[str],
        reset_to_zero: List[str],
        restore_value: List[str],
        pre_hook: Optional[Callable] = None,
        post_hook: Optional[Callable] = None,
        prune_configs_by: Optional[Dict] = None,
        warmup: Optional[int] = None,
        rep: Optional[int] = None,
        use_cuda_graph: bool = False,
    ) -> None:
        super().__init__(
            fn,
            arg_names,
            configs,
            key,
            reset_to_zero,
            restore_value,
            pre_hook,
            post_hook,
            prune_configs_by,
            warmup,
            rep,
            use_cuda_graph,
        )
        # File where we'll store results if disk caching is enabled
        self.cache_file = None
        if config_dir is not None:
            config_dir.mkdir(parents=True, exist_ok=True)
            # TODO: adjust this to also include the fn's hash?
            self.cache_file = config_dir / f"{fn.__name__}_{device_capability}.json"
            # Load any previously cached data
            if self.cache_file.exists():
                try:
                    with FILE_LOCK, open(self.cache_file, "rb") as f:
                        self.cache = {
                            k: dict_to_config(v) for k, v in json.load(f).items()
                        }
                except Exception as e:
                    # If there's some corruption or incompatibility, ignore and start fresh
                    print(
                        f"Warning: Could not load autotune cache from {self.cache_file}: {e}"
                    )
                    self.cache = {}

    def _write_cache_to_disk(self) -> None:
        """Writes the current self.cache to disk if self.cache_dir is set."""
        if self.cache_file is None:
            return
        temp_file = f"{self.cache_file}.{os.getpid()}.tmp"
        try:
            with FILE_LOCK:
                with open(temp_file, "w") as f:
                    json.dump(
                        {k: config_to_dict(v) for k, v in self.cache.items()},
                        f,
                        indent=4,
                    )
                os.replace(temp_file, self.cache_file)
        except Exception as e:
            print(f"Warning: Failed to write autotune cache: {e}")
        finally:
            if os.path.exists(temp_file):
                os.remove(temp_file)

    def run(self, *args: Any, **kwargs: Any) -> Any:
        self.nargs = dict(zip(self.arg_names, args))
        used_cached_result = True
        if len(self.configs) > 1:
            all_args = {**self.nargs, **kwargs}
            _args = {k: v for (k, v) in all_args.items() if k in self.arg_names}
            key = [_args[key] for key in self.keys if key in _args]
            for _, arg in _args.items():
                if hasattr(arg, "dtype"):
                    key.append(str(arg.dtype))
            key = "_".join(map(str, key))
            if key not in self.cache:
                # prune configs
                used_cached_result = False
                pruned_configs = self.prune_configs(kwargs)
                bench_start = time.time()
                timings = {
                    config: self._bench(*args, config=config, **kwargs)
                    for config in pruned_configs
                }
                bench_end = time.time()
                self.bench_time = bench_end - bench_start
                self.cache[key] = builtins.min(timings, key=timings.get)
                full_nargs = {**self.nargs, **kwargs, **self.cache[key].all_kwargs()}
                self.pre_hook(full_nargs, reset_only=True)
                self.configs_timings = timings

                self._write_cache_to_disk()
            config = self.cache[key]
        else:
            config = self.configs[0]
        self.best_config = config
        if os.getenv("TRITON_PRINT_AUTOTUNING", None) == "1" and not used_cached_result:
            print(
                f"Triton autotuning for function {self.base_fn.__name__} finished after "
                f"{self.bench_time:.2f}s; best config selected: {self.best_config};"
            )
        if config.pre_hook is not None:
            full_nargs = {**self.nargs, **kwargs, **config.all_kwargs()}
            config.pre_hook(full_nargs)
        ret = self.fn.run(
            *args,
            **kwargs,
            **config.all_kwargs(),
        )
        self.nargs = None
        return ret


def autotune(
    configs: List[triton.Config],
    key: List[str],
    prune_configs_by: Optional[Dict] = None,
    reset_to_zero: Optional[List[str]] = None,
    restore_value: Optional[List[str]] = None,
    pre_hook: Optional[Callable] = None,
    post_hook: Optional[Callable] = None,
    warmup: Optional[int] = None,
    rep: Optional[int] = None,
    use_cuda_graph: bool = False,
    do_bench: Optional[Callable] = None,
) -> Callable:
    """
    Decorator for auto-tuning a :code:`triton.jit`'d function.

    .. highlight:: python
    .. code-block:: python

        @triton.autotune(configs=[
            triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4),
            triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8),
          ],
          key=['x_size'] # the two above configs will be evaluated anytime
                         # the value of x_size changes
        )
        @triton.jit
        def kernel(x_ptr, x_size, **META):
            BLOCK_SIZE = META['BLOCK_SIZE']
    :note: When all the configurations are evaluated, the kernel will run multiple times.
           This means that whatever value the kernel updates will be updated multiple times.
           To avoid this undesired behavior, you can use the `reset_to_zero` argument, which
           resets the value of the provided tensor to `zero` before running any configuration.

    If the environment variable :code:`TRITON_PRINT_AUTOTUNING` is set to
    :code:`"1"`, Triton will print a message to stdout after autotuning each
    kernel, including the time spent autotuning and the best configuration.

    :param configs: a list of :code:`triton.Config` objects
    :type configs: list[triton.Config]
    :param key: a list of argument names whose change in value will trigger the evaluation of all provided configs.
    :type key: list[str]
    :param prune_configs_by: a dict of functions that are used to prune configs, fields:
        'perf_model': performance model used to predicate running time with different configs, returns running time
        'top_k': number of configs to bench
        'early_config_prune'(optional): a function used to do early prune (eg, num_stages). It takes configs:List[Config] as its input, and returns pruned configs.
    :param reset_to_zero: a list of argument names whose value will be reset to zero before evaluating any configs.
    :type reset_to_zero: list[str]
    :param restore_value: a list of argument names whose value will be restored after evaluating any configs.
    :type restore_value: list[str]
    :param pre_hook: a function that will be called before the kernel is called.
        This overrides the default pre_hook used for 'reset_to_zero' and 'restore_value'.
        'kwargs': a dict of all arguments passed to the kernel.
        'reset_only': a boolean indicating whether the pre_hook is called to reset the values only, without a corresponding post_hook.
    :type pre_hook: lambda args, reset_only
    :param post_hook: a function that will be called after the kernel is called.
        This overrides the default post_hook used for 'restore_value'.
        'kwargs': a dict of all arguments passed to the kernel.
        'exception': the exception raised by the kernel in case of a compilation or runtime error.
    :type post_hook: lambda args, exception
    :param warmup: warmup time (in ms) to pass to benchmarking (deprecated).
    :type warmup: int
    :param rep: repetition time (in ms) to pass to benchmarking (deprecated).
    :type rep: int
    :param do_bench: a benchmark function to measure the time of each run.
    :type do_bench: lambda fn, quantiles
    """

    def decorator(fn: Callable) -> Autotuner:
        return Autotuner(
            fn,
            fn.arg_names,
            configs,
            key,
            reset_to_zero,
            restore_value,
            pre_hook=pre_hook,
            post_hook=post_hook,
            prune_configs_by=prune_configs_by,
            warmup=warmup,
            rep=rep,
            use_cuda_graph=use_cuda_graph,
        )

    return decorator
