# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import json
import pathlib
import time
from typing import List, Optional, Sequence, Set, Tuple
from common.benchmark_suite import BenchmarkCase, BenchmarkSuite
from common.benchmark_config import BenchmarkConfig
from common.benchmark_definition import (
    BenchmarkInfo,
    BenchmarkResults,
    BenchmarkMetrics,
    BenchmarkRun,
    DeviceInfo,
)


class BenchmarkDriver(object):
    """Abstract driver runs the whole benchmark flow."""

    def __init__(
        self,
        device_info: DeviceInfo,
        benchmark_config: BenchmarkConfig,
        benchmark_suite: BenchmarkSuite,
        benchmark_grace_time: float = 0.0,
        verbose: bool = False,
    ):
        self.device_info = device_info
        self.config = benchmark_config
        self.benchmark_suite = benchmark_suite
        self.benchmark_grace_time = benchmark_grace_time
        self.verbose = verbose
        self.finished_benchmarks: List[Tuple[BenchmarkInfo, pathlib.Path]] = []
        self.finished_captures: List[pathlib.Path] = []
        self.benchmark_errors = []
        self._seen_benchmark_names: Set[str] = set()

    def run_benchmark_case(
        self,
        benchmark_case: BenchmarkCase,
        benchmark_results_filename: Optional[pathlib.Path],
        capture_filename: Optional[pathlib.Path],
    ) -> None:
        """Runs the benchmark case and serializes the results.

        Args:
          benchmark_case: the benchmark_case.
          benchmark_results_filename: the path to store the serialized
            BenchmarkMetrics. Benchmarking is required if set.
          capture_filename: the path to store captured trace. Trace capturing is
            required if set.

        Raises:
          Exception during benchmarking.
        """
        raise NotImplementedError("Should be overwritten by a subclass.")

    def run(self) -> None:
        """Execute the benchmark flow.

        It performs the following steps:
          1. Enumerate and filter benchmark cases.
          2. Call 'run_benchmark_case' for each benchmark case.
          3. Collect the benchmark results and captures.
        """

        self.config.benchmark_results_dir.mkdir(parents=True, exist_ok=True)
        if self.config.trace_capture_config is not None:
            self.config.trace_capture_config.capture_tmp_dir.mkdir(
                parents=True, exist_ok=True
            )

        cpu_target_arch = self.device_info.get_cpu_arch()
        gpu_target_arch = self.device_info.get_gpu_arch()
        detected_architectures = [
            arch for arch in [cpu_target_arch, gpu_target_arch] if arch is not None
        ]
        if self.config.use_compatible_filter:
            if cpu_target_arch is None:
                print(
                    "INFO: Detected unsupported CPU architecture in"
                    f' "{self.device_info}", CPU benchmarking is disabled.'
                )
            if gpu_target_arch is None:
                print(
                    "INFO: Detected unsupported GPU architecture in"
                    f' "{self.device_info}", GPU benchmarking is disabled.'
                )
            compatible_arch_filter = detected_architectures
        else:
            # No compatible filter on the target architectures.
            compatible_arch_filter = None

        drivers, loaders = self.__get_available_drivers_and_loaders()

        benchmark_cases = self.benchmark_suite.filter_benchmarks(
            available_drivers=drivers,
            available_loaders=loaders,
            target_architectures=compatible_arch_filter,
            driver_filter=self.config.driver_filter,
            mode_filter=self.config.mode_filter,
            model_name_filter=self.config.model_name_filter,
        )

        for benchmark_case in benchmark_cases:
            benchmark_info = self.__get_benchmark_info_from_case(
                benchmark_case=benchmark_case
            )
            benchmark_name = str(benchmark_info)

            if benchmark_case.target_arch not in detected_architectures:
                print(
                    f"WARNING: Benchmark '{benchmark_name}' may be incompatible"
                    f" with the detected architectures '{detected_architectures}'"
                    f" on the device. Pass --compatible-only to skip incompatible"
                    f" benchmarks."
                )

            # Sanity check for the uniqueness of benchmark names.
            if benchmark_name in self._seen_benchmark_names:
                raise ValueError(
                    f"Found duplicate benchmark {benchmark_name} in the suites."
                )
            self._seen_benchmark_names.add(benchmark_name)

            results_path, capture_path = self.__get_output_paths(benchmark_name)
            # If we continue from the previous results, check and skip if the result
            # files exist.
            if self.config.continue_from_previous:
                if results_path is not None and results_path.exists():
                    self.finished_benchmarks.append((benchmark_info, results_path))
                    results_path = None

                if capture_path is not None and capture_path.exists():
                    self.finished_captures.append(capture_path)
                    capture_path = None

            # Skip if no need to benchmark and capture.
            if results_path is None and capture_path is None:
                continue

            print(f"--> Benchmark started: {benchmark_name} <--")

            try:
                self.run_benchmark_case(benchmark_case, results_path, capture_path)
            except Exception as e:
                # Delete unfinished results if they exist.
                if results_path is not None:
                    results_path.unlink(missing_ok=True)
                if capture_path is not None:
                    capture_path.unlink(missing_ok=True)

                if not self.config.keep_going:
                    raise e

                print(f"Processing of benchmark failed with: {e}")
                self.benchmark_errors.append(e)
                continue
            finally:
                # Some grace time.
                time.sleep(self.benchmark_grace_time)

            print("Benchmark completed")

            if results_path:
                self.finished_benchmarks.append((benchmark_info, results_path))
            if capture_path:
                self.finished_captures.append(capture_path)

    def get_benchmark_results(self) -> BenchmarkResults:
        """Returns the finished benchmark results."""

        results = BenchmarkResults()
        results.set_commit(self.config.git_commit_hash)

        finished_benchmarks = sorted(
            self.finished_benchmarks, key=lambda pair: str(pair[0])
        )
        for info, path in finished_benchmarks:
            benchmark_metrics_json_object = json.loads(path.read_text())
            benchmark_run = BenchmarkRun(
                info=info,
                metrics=BenchmarkMetrics.from_json_object(
                    benchmark_metrics_json_object
                ),
            )
            results.benchmarks.append(benchmark_run)

        return results

    def get_benchmark_result_filenames(self) -> Sequence[pathlib.Path]:
        """Returns the json file paths of finished benchmarks."""
        return [path for info, path in self.finished_benchmarks]

    def get_capture_filenames(self) -> Sequence[pathlib.Path]:
        """Returns the tracy file paths of finished captures."""
        return self.finished_captures

    def get_benchmark_errors(self):
        """Returns the exceptions captured during benchmarking."""
        return self.benchmark_errors

    def __get_output_paths(self, benchmark_name: str):
        """Get output paths for the results and capture. The path of results/capture
        is None if the benchmark/capture doesn't need to be run.
        """

        benchmark_results_filename = None
        if self.config.normal_benchmark_tool_dir:
            benchmark_results_filename = (
                self.config.benchmark_results_dir / f"{benchmark_name}.json"
            )

        capture_filename = None
        if self.config.trace_capture_config:
            capture_filename = (
                self.config.trace_capture_config.capture_tmp_dir
                / f"{benchmark_name}.tracy"
            )

        return (benchmark_results_filename, capture_filename)

    def __get_benchmark_info_from_case(
        self, benchmark_case: BenchmarkCase
    ) -> BenchmarkInfo:
        run_config = benchmark_case.run_config
        run_tags = run_config.module_execution_config.tags
        gen_config = run_config.module_generation_config
        model_source = str(gen_config.imported_model.model.source_type)
        compile_tags = gen_config.compile_config.tags
        return BenchmarkInfo(
            name=run_config.name,
            model_name=benchmark_case.model_name,
            model_tags=benchmark_case.model_tags,
            model_source=model_source,
            bench_mode=run_tags,
            compile_tags=compile_tags,
            driver_info=benchmark_case.driver_info,
            device_info=self.device_info,
            run_config_id=run_config.composite_id,
        )

    def __get_available_drivers_and_loaders(
        self,
    ) -> Tuple[Sequence[str], Sequence[str]]:
        any_tool_dir = (
            self.config.normal_benchmark_tool_dir
            if self.config.normal_benchmark_tool_dir
            else self.config.trace_capture_config.traced_benchmark_tool_dir
        )
        config_txt_file_path = any_tool_dir / "build_config.txt"
        config_txt_file_lines = config_txt_file_path.read_text().splitlines()

        available_drivers = []
        available_loaders = []
        for line in config_txt_file_lines:
            name, value = line.strip().split("=")
            if value != "ON":
                continue
            if name == "IREE_HAL_DRIVER_CUDA":
                available_drivers.append("cuda")
            elif name == "IREE_HAL_DRIVER_LOCAL_SYNC":
                available_drivers.append("local-sync")
            elif name == "IREE_HAL_DRIVER_LOCAL_TASK":
                available_drivers.append("local-task")
            elif name == "IREE_HAL_DRIVER_VULKAN":
                available_drivers.append("vulkan")
            elif name == "IREE_HAL_EXECUTABLE_LOADER_EMBEDDED_ELF":
                available_loaders.append("embedded-elf")
            elif name == "IREE_HAL_EXECUTABLE_LOADER_SYSTEM_LIBRARY":
                available_loaders.append("system-library")
            elif name == "IREE_HAL_EXECUTABLE_LOADER_VMVX_MODULE":
                available_loaders.append("vmvx-module")
            else:
                continue

        if self.verbose:
            available_drivers_str = ", ".join(available_drivers)
            print(f"Available drivers: {available_drivers_str}")
            available_loaders_str = ", ".join(available_loaders)
            print(f"Available loaders: {available_loaders_str}")

        return available_drivers, available_loaders