#!/usr/bin/env python3
# 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
"""Collect compilation statistics from benchmark suites.

See docs/developers/developing_iree/benchmark_suites.md for how to build the
benchmark suites.
"""

import pathlib
import sys

# Add build_tools python dir to the search path.
sys.path.insert(0, str(pathlib.Path(__file__).parent.with_name("python")))

import argparse
import json
import os
import re
import zipfile

from dataclasses import asdict, dataclass
from typing import BinaryIO, Dict, List, Optional, TextIO

from common import benchmark_definition
from common.benchmark_definition import (
    CompilationInfo,
    CompilationResults,
    CompilationStatistics,
    ModuleComponentSizes,
    get_git_commit_hash,
)
from common import benchmark_config
from e2e_test_artifacts import iree_artifacts
from e2e_test_framework import serialization
from e2e_test_framework.definitions import iree_definitions

BENCHMARK_FLAGFILE = "flagfile"
MODULE_DIR = "vmfb"
MODULE_FILE_EXTENSION = "vmfb"
NINJA_LOG_HEADER = "ninja log v5"
NINJA_BUILD_LOG = ".ninja_log"
COMPILATION_STATS_MODULE_SUFFIX = "compile-stats"
E2E_TEST_ARTIFACTS_REL_PATH = "e2e_test_artifacts"

VM_COMPONENT_NAME = "module.fb"
CONST_COMPONENT_NAME = "_const.bin"
DISPATCH_COMPONENT_PATTERNS = [
    r".+_embedded_elf_.+\.so",
    r".+_vulkan_spirv_fb\.fb",
    r".+_cuda_nvptx_fb\.fb",
    r".+_vmvx_bytecode_fb\.fb",
]


@dataclass(frozen=True)
class ModuleInfo(object):
    module_path: pathlib.Path
    stream_stats_path: pathlib.Path


def match_module_cmake_target(module_path: pathlib.PurePath) -> Optional[str]:
    if module_path.match(
        f"{E2E_TEST_ARTIFACTS_REL_PATH}/iree_*/" f"{iree_artifacts.MODULE_FILENAME}"
    ):
        # <e2e test artifacts dir>/iree_<module dir>/<module filename>
        path_parts = module_path.parts[-3:]
        # Join to get the CMake target name. This is *not* a filesystem path, so we
        # don't want \ separators on Windows that we would get with os.path.join().
        return "/".join(path_parts)

    return None


def parse_compilation_time_from_ninja_log(log: TextIO) -> Dict[str, int]:
    """Retrieve the compilation time (ms) from the Ninja build log.

    Returns:
      Map of target name and compilation time in ms.
    """

    target_build_time_map = {}
    header = log.readline()
    if NINJA_LOG_HEADER not in header:
        raise NotImplementedError(f"Unsupported ninja log version: {header}")

    for line in log:
        start_time, end_time, _, target, _ = line.strip().split("\t")
        cmake_target = match_module_cmake_target(pathlib.PurePath(target))
        if cmake_target is None:
            continue

        start_time = int(start_time)
        end_time = int(end_time)
        target_build_time_map[cmake_target] = end_time - start_time

    return target_build_time_map


def get_module_component_info(
    module: BinaryIO, module_file_bytes: int
) -> ModuleComponentSizes:
    with zipfile.ZipFile(module) as module_zipfile:
        size_map = dict(
            (info.filename, info.file_size) for info in module_zipfile.infolist()
        )

    identified_names = set()
    if VM_COMPONENT_NAME in size_map:
        vm_component_bytes = size_map[VM_COMPONENT_NAME]
        identified_names.add(VM_COMPONENT_NAME)
    else:
        vm_component_bytes = 0

    if CONST_COMPONENT_NAME in size_map:
        const_component_bytes = size_map[CONST_COMPONENT_NAME]
        identified_names.add(CONST_COMPONENT_NAME)
    else:
        const_component_bytes = 0

    total_dispatch_component_bytes = 0
    for filename, size in size_map.items():
        for pattern in DISPATCH_COMPONENT_PATTERNS:
            if re.match(pattern, filename):
                total_dispatch_component_bytes += size
                identified_names.add(filename)
                break

    if identified_names != set(size_map.keys()):
        raise RuntimeError(f"Unrecognized components in the module: {size_map.keys()}.")

    return ModuleComponentSizes(
        file_bytes=module_file_bytes,
        vm_component_bytes=vm_component_bytes,
        const_component_bytes=const_component_bytes,
        total_dispatch_component_bytes=total_dispatch_component_bytes,
    )


def get_module_map_from_compilation_benchmark_config(
    compilation_benchmark_config_data: TextIO, e2e_test_artifacts_dir: pathlib.PurePath
) -> Dict[CompilationInfo, ModuleInfo]:
    benchmark_config = json.load(compilation_benchmark_config_data)
    gen_configs = serialization.unpack_and_deserialize(
        data=benchmark_config["generation_configs"],
        root_type=List[iree_definitions.ModuleGenerationConfig],
    )
    module_map = {}
    for gen_config in gen_configs:
        model = gen_config.imported_model.model
        compile_config = gen_config.compile_config
        target_archs = []
        for compile_target in compile_config.compile_targets:
            arch = compile_target.target_architecture
            target_archs.append(
                (
                    f"{arch.type.value}-{arch.architecture}-{arch.microarchitecture}-"
                    f"{compile_target.target_abi.value}"
                )
            )
        compilation_info = CompilationInfo(
            name=gen_config.name,
            model_name=model.name,
            model_tags=tuple(model.tags),
            model_source=model.source_type.value,
            target_arch=f"[{','.join(target_archs)}]",
            compile_tags=tuple(compile_config.tags),
            gen_config_id=gen_config.composite_id,
        )
        module_dir_path = pathlib.Path(
            iree_artifacts.get_module_dir_path(
                module_generation_config=gen_config, root_path=e2e_test_artifacts_dir
            )
        )
        module_path = module_dir_path / iree_artifacts.MODULE_FILENAME
        stream_stats_path = module_dir_path / iree_artifacts.SCHEDULING_STATS_FILENAME
        module_map[compilation_info] = ModuleInfo(
            module_path=module_path, stream_stats_path=stream_stats_path
        )

    return module_map


def _check_dir_path(path_str: str) -> pathlib.Path:
    path = pathlib.Path(path_str)
    if not path.is_dir():
        raise argparse.ArgumentTypeError(f"{path} is not a directory.")
    return path


def _check_file_path(path_str: str) -> pathlib.Path:
    path = pathlib.Path(path_str)
    if not path.is_file():
        raise argparse.ArgumentTypeError(f"{path} is not a file.")
    return path


def _parse_arguments():
    """Returns an argument parser with common options."""

    parser = argparse.ArgumentParser(
        description="Collect compilation statistics from benchmark suites."
    )
    parser.add_argument(
        "--compilation_benchmark_config",
        type=_check_file_path,
        required=True,
        help="Exported compilation benchmark config of e2e test artifacts.",
    )
    parser.add_argument(
        "--build_log",
        type=_check_file_path,
        required=True,
        help="Path to the ninja build log.",
    )
    parser.add_argument(
        "--e2e_test_artifacts_dir",
        type=_check_dir_path,
        required=True,
        help="Path to the e2e test artifacts directory.",
    )
    parser.add_argument("--output", type=pathlib.Path, help="Path to output JSON file.")

    return parser.parse_args()


def main(args: argparse.Namespace):
    config_data = args.compilation_benchmark_config.open("r")
    module_map = get_module_map_from_compilation_benchmark_config(
        compilation_benchmark_config_data=config_data,
        e2e_test_artifacts_dir=args.e2e_test_artifacts_dir,
    )
    build_log_path = args.build_log

    with build_log_path.open("r") as log_file:
        target_build_time_map = parse_compilation_time_from_ninja_log(log_file)

    compilation_statistics_list = []
    for compilation_info, module_info in module_map.items():
        module_path = module_info.module_path
        with module_path.open("rb") as module_file:
            module_component_sizes = get_module_component_info(
                module_file, module_path.stat().st_size
            )

        cmake_target = match_module_cmake_target(module_path)
        if cmake_target is None:
            raise RuntimeError(
                f"Module path isn't a module cmake target: {module_path}"
            )
        compilation_time_ms = target_build_time_map[cmake_target]

        stream_stats_json = json.loads(module_info.stream_stats_path.read_text())
        exec_stats_json = stream_stats_json["stream-aggregate"]["execution"]
        ir_stats = benchmark_definition.IRStatistics(
            stream_dispatch_count=exec_stats_json["dispatch-count"]
        )

        compilation_statistics = CompilationStatistics(
            compilation_info=compilation_info,
            module_component_sizes=module_component_sizes,
            compilation_time_ms=compilation_time_ms,
            ir_stats=ir_stats,
        )
        compilation_statistics_list.append(compilation_statistics)

    commit = get_git_commit_hash("HEAD")
    compilation_results = CompilationResults(
        commit=commit, compilation_statistics=compilation_statistics_list
    )

    json_output = json.dumps(asdict(compilation_results), indent=2)
    if args.output is None:
        print(json_output)
    else:
        args.output.write_text(json_output)


if __name__ == "__main__":
    main(_parse_arguments())