#!/usr/bin/env python3
# Copyright 2021 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
"""Generates benchmark results as pull request comments.

This script is meant to be used by CI and uses pip package "markdown_strings".
"""

import sys
import pathlib

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

import argparse
import dataclasses
import json
from typing import Any, Dict, Optional, Set, Tuple

import markdown_strings as md
import requests

from common import benchmark_definition, benchmark_presentation, common_arguments
from reporting import benchmark_comment

GITHUB_IREE_REPO_PREFIX = "https://github.com/openxla/iree"
IREE_DASHBOARD_URL = "https://perf.iree.dev/apis/v2"
IREE_PROJECT_ID = "IREE"
# The maximal numbers of trials when querying base commit benchmark results.
MAX_BASE_COMMIT_QUERY_COUNT = 10
# The max number of rows to show per table.
TABLE_SIZE_CUT = 3
THIS_DIRECTORY = pathlib.Path(__file__).resolve().parent


@dataclasses.dataclass(frozen=True)
class CommentDef(object):
    title: str
    type_id: str


# Map from comment type to comment definition.
COMMENT_DEF_MAP = {
    "android-benchmark-summary": CommentDef(
        title="Abbreviated Android Benchmark Summary",
        type_id="bf8cdf94-a992-466d-b11c-778cbd805a22",
    ),
    "linux-benchmark-summary": CommentDef(
        title="Abbreviated Linux Benchmark Summary",
        type_id="37549014-3c67-4e74-8d88-8e929231abe3",
    ),
    "benchmark-summary": CommentDef(
        title="Abbreviated Benchmark Summary",
        type_id="5b42cbfe-26a0-4164-a51c-07f06762e2dc",
    ),
}


def get_git_total_commit_count(commit: str, verbose: bool = False) -> int:
    """Gets the total commit count in history ending with the given commit."""
    # TODO(#11703): Should use --first-parent here. See issue for the required
    # work.
    count = benchmark_definition.execute_cmd_and_get_stdout(
        ["git", "rev-list", "--count", commit], cwd=THIS_DIRECTORY, verbose=verbose
    )
    return int(count)


def get_from_dashboard(
    url: str, payload: Dict[str, Any], verbose: bool = False
) -> Dict[str, Dict[str, Any]]:
    headers = {"Content-type": "application/json"}
    data = json.dumps(payload)

    if verbose:
        print(f"API request payload: {data}")

    response = requests.get(url, data=data, headers=headers)
    code = response.status_code
    if code != 200:
        raise requests.RequestException(
            f"Failed to get from dashboard server with status code {code}"
        )

    data = response.json()
    if verbose:
        print(f"Queried base benchmark data: {data}")
    return data


BenchmarkQueryResults = Dict[str, Dict[str, Any]]


def query_base_benchmark_results(
    commit: str, verbose: bool = False
) -> BenchmarkQueryResults:
    """Queries the benchmark results for the given commit."""
    build_id = get_git_total_commit_count(commit, verbose)
    payload = {"projectId": IREE_PROJECT_ID, "buildId": build_id}
    return get_from_dashboard(
        f"{IREE_DASHBOARD_URL}/getBuild", payload, verbose=verbose
    )


@dataclasses.dataclass(frozen=True)
class ComparableBenchmarkResults(object):
    commit_sha: str
    benchmark_results: BenchmarkQueryResults


def _find_comparable_benchmark_results(
    start_commit: str, required_benchmark_keys: Set[str], verbose: bool = False
) -> Optional[ComparableBenchmarkResults]:
    cmds = [
        "git",
        "rev-list",
        "--first-parent",
        f"--max-count={MAX_BASE_COMMIT_QUERY_COUNT}",
        start_commit,
    ]
    output = benchmark_definition.execute_cmd_and_get_stdout(
        cmds, cwd=THIS_DIRECTORY, verbose=verbose
    )
    previous_commits = output.splitlines()
    # Try to query some base benchmark to diff against, from the top of the
    # tree. Bail out if the maximal trial number is exceeded.
    for base_commit in previous_commits:
        base_benchmarks = query_base_benchmark_results(
            commit=base_commit, verbose=verbose
        )
        base_benchmark_keys = set(base_benchmarks.keys())
        if required_benchmark_keys <= base_benchmark_keys:
            return ComparableBenchmarkResults(
                commit_sha=base_commit, benchmark_results=base_benchmarks
            )

    return None


def _get_git_commit_hash(ref: str, verbose: bool = False) -> str:
    """Gets the commit hash for the given commit."""
    return benchmark_definition.execute_cmd_and_get_stdout(
        ["git", "rev-parse", ref], cwd=THIS_DIRECTORY, verbose=verbose
    )


def _get_git_merge_base_commit(
    pr_commit: str, target_branch: str, verbose: bool = False
) -> str:
    return benchmark_definition.execute_cmd_and_get_stdout(
        args=["git", "merge-base", target_branch, pr_commit],
        cwd=THIS_DIRECTORY,
        verbose=verbose,
    )


def _get_benchmark_result_markdown(
    execution_benchmarks: Dict[str, benchmark_presentation.AggregateBenchmarkLatency],
    compilation_metrics: Dict[str, benchmark_presentation.CompilationMetrics],
    pr_url: str,
    build_url: str,
    comment_def: CommentDef,
    commit_info_md: str,
) -> Tuple[str, str]:
    """Gets the full/abbreviated markdown summary of all benchmarks in files."""

    pr_info = md.link("Pull request", pr_url)
    build_info = md.link("Build", build_url)

    # Compose the full benchmark tables.
    full_table = [md.header("Full Benchmark Summary", 2)]
    full_table.append(md.unordered_list([commit_info_md, pr_info, build_info]))

    # Compose the abbreviated benchmark tables.
    abbr_table = [md.header(comment_def.title, 2)]
    abbr_table.append(commit_info_md)

    if len(execution_benchmarks) > 0:
        full_table.append(
            benchmark_presentation.categorize_benchmarks_into_tables(
                execution_benchmarks
            )
        )

        abbr_benchmarks_tables = (
            benchmark_presentation.categorize_benchmarks_into_tables(
                execution_benchmarks, TABLE_SIZE_CUT
            )
        )
        if len(abbr_benchmarks_tables) == 0:
            abbr_table.append("No improved or regressed benchmarks 🏖️")
        else:
            abbr_table.append(abbr_benchmarks_tables)

    # Compose the full compilation metrics tables.
    if len(compilation_metrics) > 0:
        full_table.append(
            benchmark_presentation.categorize_compilation_metrics_into_tables(
                compilation_metrics
            )
        )

        abbr_compilation_metrics_tables = (
            benchmark_presentation.categorize_compilation_metrics_into_tables(
                compilation_metrics, TABLE_SIZE_CUT
            )
        )
        if len(abbr_compilation_metrics_tables) == 0:
            abbr_table.append("No improved or regressed compilation metrics 🏖️")
        else:
            abbr_table.append(abbr_compilation_metrics_tables)

    abbr_table.append("For more information:")
    # We don't know until a Gist is really created. Use a placeholder for now and
    # replace later.
    full_result_info = md.link(
        "Full benchmark result tables", benchmark_comment.GIST_LINK_PLACEHORDER
    )
    abbr_table.append(md.unordered_list([full_result_info, build_info]))

    # Append the unique comment type id to help identify and update the existing
    # comment.
    abbr_table.append(f"<!--Comment type id: {comment_def.type_id}-->")

    return "\n\n".join(full_table), "\n\n".join(abbr_table)


def parse_arguments():
    """Parses command-line options."""

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--benchmark_files",
        metavar="<benchmark-json-files>",
        default=[],
        action="append",
        help=(
            "Paths to the JSON files containing benchmark results, " "accepts wildcards"
        ),
    )
    parser.add_argument(
        "--compile_stats_files",
        metavar="<compile-stats-json-files>",
        default=[],
        action="append",
        help=(
            "Paths to the JSON files containing compilation statistics, "
            "accepts wildcards"
        ),
    )
    parser.add_argument("--pr_number", required=True, type=int, help="PR number")
    parser.add_argument(
        "--pr_committish", type=str, default="HEAD", help="PR commit hash or ref"
    )
    parser.add_argument(
        "--pr_base_branch", type=str, default=None, help="Base branch to merge the PR."
    )
    parser.add_argument(
        "--comment_type",
        required=True,
        choices=COMMENT_DEF_MAP.keys(),
        help="Type of summary comment",
    )
    parser.add_argument(
        "--build_url",
        required=True,
        type=str,
        help="CI build page url to show in the report",
    )
    parser.add_argument("--output", type=pathlib.Path, default=None)
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print internal information during execution",
    )

    return parser.parse_args()


def main(args):
    benchmark_files = common_arguments.expand_and_check_file_paths(args.benchmark_files)
    compile_stats_files = common_arguments.expand_and_check_file_paths(
        args.compile_stats_files
    )

    pr_commit = _get_git_commit_hash(ref=args.pr_committish, verbose=args.verbose)
    execution_benchmarks = benchmark_presentation.aggregate_all_benchmarks(
        benchmark_files=benchmark_files, expected_pr_commit=pr_commit
    )
    compilation_metrics = benchmark_presentation.collect_all_compilation_metrics(
        compile_stats_files=compile_stats_files, expected_pr_commit=pr_commit
    )

    if args.pr_base_branch is None:
        pr_base_commit = None
    else:
        pr_base_commit = _get_git_merge_base_commit(
            pr_commit=pr_commit, target_branch=args.pr_base_branch, verbose=args.verbose
        )

    if pr_base_commit is None:
        comparable_results = None
    else:
        required_benchmark_keys = set(execution_benchmarks.keys())
        for target_id in compilation_metrics:
            for mapper in benchmark_presentation.COMPILATION_METRICS_TO_TABLE_MAPPERS:
                required_benchmark_keys.add(mapper.get_series_id(target_id))

        comparable_results = _find_comparable_benchmark_results(
            start_commit=pr_base_commit,
            required_benchmark_keys=required_benchmark_keys,
            verbose=args.verbose,
        )

    if comparable_results is None:
        comparable_commit = None
    else:
        comparable_commit = comparable_results.commit_sha
        # Update the execution benchmarks with base numbers.
        for bench in execution_benchmarks:
            base_benchmark = comparable_results.benchmark_results[bench]
            if base_benchmark["sampleUnit"] != "ns":
                raise ValueError("Only support nanoseconds for latency sample.")
            execution_benchmarks[bench].base_mean_time = base_benchmark["sample"]

        # Update the compilation metrics with base numbers.
        for target_id, metrics in compilation_metrics.items():
            updated_metrics = metrics
            for mapper in benchmark_presentation.COMPILATION_METRICS_TO_TABLE_MAPPERS:
                base_benchmark = comparable_results.benchmark_results[
                    mapper.get_series_id(target_id)
                ]
                if base_benchmark["sampleUnit"] != mapper.get_unit():
                    raise ValueError("Unit of the queried sample is mismatched.")
                updated_metrics = mapper.update_base_value(
                    updated_metrics, base_benchmark["sample"]
                )
            compilation_metrics[target_id] = updated_metrics

    pr_commit_link = md.link(pr_commit, f"{GITHUB_IREE_REPO_PREFIX}/commit/{pr_commit}")
    commit_info_md = f"@ commit {pr_commit_link}"
    if comparable_commit is not None:
        baseline_commit_link = md.link(
            comparable_commit, f"{GITHUB_IREE_REPO_PREFIX}/commit/{comparable_commit}"
        )
        commit_info_md += f" (vs. base {baseline_commit_link})"
    elif pr_base_commit is not None:
        commit_info_md += " (no previous benchmark results to compare)"

    comment_def = COMMENT_DEF_MAP[args.comment_type]
    full_md, abbr_md = _get_benchmark_result_markdown(
        execution_benchmarks=execution_benchmarks,
        compilation_metrics=compilation_metrics,
        pr_url=f"{GITHUB_IREE_REPO_PREFIX}/pull/{args.pr_number}",
        build_url=args.build_url,
        comment_def=comment_def,
        commit_info_md=commit_info_md,
    )

    comment_data = benchmark_comment.CommentData(
        type_id=comment_def.type_id,
        abbr_md=abbr_md,
        full_md=full_md,
        unverified_pr_number=args.pr_number,
    )
    comment_json_data = json.dumps(dataclasses.asdict(comment_data), indent=2)
    if args.output is None:
        print(comment_json_data)
    else:
        args.output.write_text(comment_json_data)


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