#!/usr/bin/env python3
# Copyright 2023 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
"""Miscellaneous tool to help work with benchmark suite and benchmark CI."""

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 shlex
import subprocess
from typing import List, Optional, Sequence

from e2e_test_artifacts import model_artifacts, iree_artifacts
from e2e_test_framework import serialization
from e2e_test_framework.definitions import iree_definitions

IREE_COMPILER_NAME = "iree-compile"


def _convert_to_cmd_string(cmds: Sequence[str]) -> str:
  if os.name == "nt":
    # list2cmdline is an undocumented method for Windows command lines. Python
    # doesn't provide an official method for quoting Windows command lines and
    # the correct implementation is slightly non-trivial. Use the undocumented
    # method for now and can be rewritten with our own implementation later.
    # See https://learn.microsoft.com/en-us/archive/blogs/twistylittlepassagesallalike/everyone-quotes-command-line-arguments-the-wrong-way
    return subprocess.list2cmdline(cmds)

  return " ".join(shlex.quote(cmd) for cmd in cmds)


def _dump_cmds_of_generation_config(
    gen_config: iree_definitions.ModuleGenerationConfig,
    root_path: pathlib.PurePath = pathlib.PurePath()):

  imported_model = gen_config.imported_model
  imported_model_path = iree_artifacts.get_imported_model_path(
      imported_model=imported_model, root_path=root_path)
  module_dir_path = iree_artifacts.get_module_dir_path(
      module_generation_config=gen_config, root_path=root_path)
  module_path = module_dir_path / iree_artifacts.MODULE_FILENAME
  compile_cmds = [
      IREE_COMPILER_NAME,
      str(imported_model_path), "-o",
      str(module_path)
  ]
  compile_cmds += gen_config.materialize_compile_flags(
      module_dir_path=module_dir_path)
  compile_cmd_str = _convert_to_cmd_string(compile_cmds)

  if imported_model.import_config.tool == iree_definitions.ImportTool.NONE:
    import_cmd_str = "# (Source model is already in MLIR)"
  else:
    source_model_path = model_artifacts.get_model_path(
        model=imported_model.model, root_path=root_path)
    import_cmds = [
        imported_model.import_config.tool.value,
        str(source_model_path), "-o",
        str(imported_model_path)
    ]
    import_cmds += imported_model.import_config.materialize_import_flags(
        model=imported_model.model)
    import_cmd_str = _convert_to_cmd_string(import_cmds)

  # Insert a blank line after each command to help read with line wrap.
  return [
      "Compile Module:", compile_cmd_str, "", "Import Model:", import_cmd_str,
      ""
  ]


def _dump_cmds_from_run_config(
    run_config: iree_definitions.E2EModelRunConfig,
    root_path: pathlib.PurePath = pathlib.PurePath()):

  gen_config = run_config.module_generation_config
  module_path = iree_artifacts.get_module_dir_path(
      module_generation_config=gen_config,
      root_path=root_path) / iree_artifacts.MODULE_FILENAME

  run_cmds = [run_config.tool.value, f"--module={module_path}"]
  run_cmds += run_config.materialize_run_flags()
  # Insert a blank line after the command to help read with line wrap.
  lines = ["Run Module:", _convert_to_cmd_string(run_cmds), ""]
  lines += _dump_cmds_of_generation_config(gen_config=gen_config,
                                           root_path=root_path)
  return lines


def _dump_cmds_handler(e2e_test_artifacts_dir: pathlib.Path,
                       execution_benchmark_config: Optional[pathlib.Path],
                       compilation_benchmark_config: Optional[pathlib.Path],
                       benchmark_id: Optional[str], **_unused_args):
  lines = []

  if execution_benchmark_config is not None:
    benchmark_groups = json.loads(execution_benchmark_config.read_text())
    for target_device, benchmark_group in benchmark_groups.items():
      run_configs = serialization.unpack_and_deserialize(
          data=benchmark_group["run_configs"],
          root_type=List[iree_definitions.E2EModelRunConfig])
      for run_config in run_configs:
        if benchmark_id is not None and benchmark_id != run_config.composite_id:
          continue

        lines.append("################")
        lines.append("")
        lines.append(f"Execution Benchmark ID: {run_config.composite_id}")
        lines.append(f"Name: {run_config}")
        lines.append(f"Target Device: {target_device}")
        lines.append("")
        lines += _dump_cmds_from_run_config(run_config=run_config,
                                            root_path=e2e_test_artifacts_dir)

  if compilation_benchmark_config is not None:
    benchmark_config = json.loads(compilation_benchmark_config.read_text())
    gen_configs = serialization.unpack_and_deserialize(
        data=benchmark_config["generation_configs"],
        root_type=List[iree_definitions.ModuleGenerationConfig])
    for gen_config in gen_configs:
      if benchmark_id is not None and benchmark_id != gen_config.composite_id:
        continue

      lines.append("################")
      lines.append("")
      lines.append(f"Compilation Benchmark ID: {gen_config.composite_id}")
      lines.append(f"Name: {gen_config}")
      lines.append("")
      lines += _dump_cmds_of_generation_config(gen_config=gen_config,
                                               root_path=e2e_test_artifacts_dir)

  print(*lines, sep="\n")


def _parse_arguments() -> argparse.Namespace:
  parser = argparse.ArgumentParser(
      description=
      "Miscellaneous tool to help work with benchmark suite and benchmark CI.")

  subparser = parser.add_subparsers(required=True, title="operation")
  dump_cmds_parser = subparser.add_parser(
      "dump-cmds",
      help="Dump the commands to compile and run benchmarks manually.")
  dump_cmds_parser.add_argument(
      "--e2e_test_artifacts_dir",
      type=pathlib.PurePath,
      default=pathlib.Path(),
      help="E2E test artifacts root path used in the outputs of artifact paths")
  dump_cmds_parser.add_argument("--benchmark_id",
                                type=str,
                                help="Only dump the benchmark with this id")
  dump_cmds_parser.add_argument(
      "--execution_benchmark_config",
      type=pathlib.Path,
      help="Config file exported from export_benchmark_config.py execution")
  dump_cmds_parser.add_argument(
      "--compilation_benchmark_config",
      type=pathlib.Path,
      help="Config file exported from export_benchmark_config.py compilation")
  dump_cmds_parser.set_defaults(handler=_dump_cmds_handler)

  args = parser.parse_args()
  if (args.execution_benchmark_config is None and
      args.compilation_benchmark_config is None):
    parser.error("At least one of --execution_benchmark_config or "
                 "--compilation_benchmark_config must be set.")

  return args


def main(args: argparse.Namespace):
  args.handler(**vars(args))


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