#!/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
"""Determines whether CI should run on a given PR.

The following environment variables are required:
- GITHUB_EVENT_NAME: GitHub event name, e.g. pull_request.
- GITHUB_OUTPUT: path to write workflow output variables.
- GITHUB_STEP_SUMMARY: path to write workflow summary output.

When GITHUB_EVENT_NAME is "pull_request", there are additional environment
variables to be set:
- PR_BRANCH (required): PR source branch.
- PR_TITLE (required): PR title.
- PR_BODY (optional): PR description.
- PR_LABELS (optional): JSON list of PR label names.
- BASE_REF (required): base commit SHA of the PR.
- ORIGINAL_PR_TITLE (optional): PR title from the original PR event, showing a
    notice if PR_TITLE is different.
- ORIGINAL_PR_BODY (optional): PR description from the original PR event,
    showing a notice if PR_BODY is different. ORIGINAL_PR_TITLE must also be
    set.
- ORIGINAL_PR_LABELS (optional): PR labels from the original PR event, showing a
    notice if PR_LABELS is different. ORIGINAL_PR_TITLE must also be set.

Exit code 0 indicates that it should and exit code 2 indicates that it should
not.
"""

import difflib
import fnmatch
import json
import os
import re
import subprocess
import textwrap
from typing import Iterable, List, Mapping, Sequence, Tuple

SKIP_CI_KEY = "skip-ci"
RUNNER_ENV_KEY = "runner-env"
BENCHMARK_EXTRA_KEY = "benchmark-extra"
# Trailer to prevent benchmarks from always running on LLVM integration PRs.
SKIP_LLVM_INTEGRATE_BENCHMARK_KEY = "skip-llvm-integrate-benchmark"

# Note that these are fnmatch patterns, which are not the same as gitignore
# patterns because they don't treat '/' specially. The standard library doesn't
# contain a function for gitignore style "wildmatch". There's a third-party
# library pathspec (https://pypi.org/project/pathspec/), but it doesn't seem
# worth the dependency.
SKIP_PATH_PATTERNS = [
    "docs/*",
    "third_party/mkdocs-material/*",
    "experimental/*",
    "build_tools/buildkite/*",
    # These configure the runners themselves and don't affect presubmit.
    "build_tools/github_actions/runner/*",
    ".github/ISSUE_TEMPLATE/*",
    "*.cff",
    "*.clang-format",
    "*.git-ignore",
    "*.md",
    "*.natvis",
    "*.pylintrc",
    "*.rst",
    "*.toml",
    "*.yamllint.yml",
    "*.yapf",
    "*CODEOWNERS",
    "*AUTHORS",
    "*LICENSE",
]

RUNNER_ENV_DEFAULT = "prod"
RUNNER_ENV_OPTIONS = [RUNNER_ENV_DEFAULT, "testing"]

DEFAULT_BENCHMARK_PRESET_GROUP = [
    "cuda",
    "x86_64",
    "android-cpu",
    "android-gpu",
    "vulkan-nvidia",
    "comp-stats",
]
DEFAULT_BENCHMARK_PRESET = "default"
LARGE_BENCHMARK_PRESET_GROUP = ["cuda-large", "x86_64-large"]
# All available benchmark preset options including experimental presets.
BENCHMARK_PRESET_OPTIONS = DEFAULT_BENCHMARK_PRESET_GROUP + LARGE_BENCHMARK_PRESET_GROUP
BENCHMARK_LABEL_PREFIX = "benchmarks"

PR_DESCRIPTION_TEMPLATE = "{title}" "\n\n" "{body}"

# Patterns to detect "LLVM integration" PRs, i.e. changes that update the
# third_party/llvm-project submodule. This should only include PRs
# intended to be merged and should exclude test/draft PRs as well as
# PRs that include temporary patches to the submodule during review.
# See also: https://github.com/openxla/iree/issues/12268
LLVM_INTEGRATE_TITLE_PATTERN = re.compile("^integrate.+llvm-project", re.IGNORECASE)
LLVM_INTEGRATE_BRANCH_PATTERN = re.compile("bump-llvm|llvm-bump", re.IGNORECASE)
LLVM_INTEGRATE_LABEL = "llvm-integrate"


def skip_path(path: str) -> bool:
    return any(fnmatch.fnmatch(path, pattern) for pattern in SKIP_PATH_PATTERNS)


def set_output(d: Mapping[str, str]):
    print(f"Setting outputs: {d}")
    step_output_file = os.environ["GITHUB_OUTPUT"]
    with open(step_output_file, "a") as f:
        f.writelines(f"{k}={v}" + "\n" for k, v in d.items())


def write_job_summary(summary: str):
    """Write markdown messages on Github workflow UI.
    See https://docs.github.com/en/actions/using-workflows/workflow-commands-for-github-actions#adding-a-job-summary
    """
    step_summary_file = os.environ["GITHUB_STEP_SUMMARY"]
    with open(step_summary_file, "a") as f:
        # Use double newlines to split sections in markdown.
        f.write(summary + "\n\n")


def check_description_and_show_diff(
    original_description: str,
    original_labels: Sequence[str],
    current_description: str,
    current_labels: Sequence[str],
):
    original_labels = sorted(original_labels)
    current_labels = sorted(current_labels)
    if (
        original_description == current_description
        and original_labels == current_labels
    ):
        return

    description_diffs = difflib.unified_diff(
        original_description.splitlines(keepends=True),
        current_description.splitlines(keepends=True),
    )
    description_diffs = "".join(description_diffs)

    if description_diffs != "":
        description_diffs = textwrap.dedent(
            """\
    ```diff
    {}
    ```
    """
        ).format(description_diffs)

    if original_labels == current_labels:
        label_diffs = ""
    else:
        label_diffs = textwrap.dedent(
            """\
    ```
    Original labels: {original_labels}
    Current labels: {current_labels}
    ```
    """
        ).format(original_labels=original_labels, current_labels=current_labels)

    write_job_summary(
        textwrap.dedent(
            """\
  :pushpin: Using the PR description and labels different from the original PR event that started this workflow.

  <details>
  <summary>Click to show diff (original vs. current)</summary>

  {description_diffs}

  {label_diffs}
  </details>"""
        ).format(description_diffs=description_diffs, label_diffs=label_diffs)
    )


def get_trailers_and_labels(is_pr: bool) -> Tuple[Mapping[str, str], List[str]]:
    if not is_pr:
        return ({}, [])

    title = os.environ["PR_TITLE"]
    body = os.environ.get("PR_BODY", "")
    labels = json.loads(os.environ.get("PR_LABELS", "[]"))
    original_title = os.environ.get("ORIGINAL_PR_TITLE")
    original_body = os.environ.get("ORIGINAL_PR_BODY", "")
    original_labels = json.loads(os.environ.get("ORIGINAL_PR_LABELS", "[]"))

    description = PR_DESCRIPTION_TEMPLATE.format(title=title, body=body)

    # PR information can be fetched from API for the latest updates. If
    # ORIGINAL_PR_TITLE is set, compare the current and original description and
    # show a notice if they are different. This is mostly to inform users that the
    # workflow might not parse the PR description they expect.
    if original_title is not None:
        original_description = PR_DESCRIPTION_TEMPLATE.format(
            title=original_title, body=original_body
        )
        print(
            "Original PR description and labels:",
            original_description,
            original_labels,
            sep="\n",
        )
        check_description_and_show_diff(
            original_description=original_description,
            original_labels=original_labels,
            current_description=description,
            current_labels=labels,
        )

    print("Parsing PR description and labels:", description, labels, sep="\n")

    trailer_lines = subprocess.run(
        ["git", "interpret-trailers", "--parse", "--no-divider"],
        input=description,
        stdout=subprocess.PIPE,
        check=True,
        text=True,
        timeout=60,
    ).stdout.splitlines()
    trailer_map = {
        k.lower().strip(): v.strip()
        for k, v in (line.split(":", maxsplit=1) for line in trailer_lines)
    }
    return (trailer_map, labels)


def get_modified_paths(base_ref: str) -> Iterable[str]:
    return subprocess.run(
        ["git", "diff", "--name-only", base_ref],
        stdout=subprocess.PIPE,
        check=True,
        text=True,
        timeout=60,
    ).stdout.splitlines()


def modifies_included_path(base_ref: str) -> bool:
    return any(not skip_path(p) for p in get_modified_paths(base_ref))


def should_run_ci(is_pr: bool, trailers: Mapping[str, str]) -> bool:
    if not is_pr:
        print(
            "Running CI independent of diff because run was not triggered by a"
            " pull request event."
        )
        return True

    if SKIP_CI_KEY in trailers:
        print(f"Not running CI because PR description has '{SKIP_CI_KEY}' trailer.")
        return False

    base_ref = os.environ["BASE_REF"]
    try:
        modifies = modifies_included_path(base_ref)
    except TimeoutError as e:
        print("Computing modified files timed out. Running the CI")
        return True

    if not modifies:
        print("Skipping CI because all modified files are marked as excluded.")
        return False

    print("CI should run")
    return True


def get_runner_env(trailers: Mapping[str, str]) -> str:
    runner_env = trailers.get(RUNNER_ENV_KEY)
    if runner_env is None:
        print(
            f"Using '{RUNNER_ENV_DEFAULT}' runners because '{RUNNER_ENV_KEY}'"
            f" not found in {trailers}"
        )
        runner_env = RUNNER_ENV_DEFAULT
    else:
        print(f"Using runner environment '{runner_env}' from PR description trailers")
    return runner_env


def get_benchmark_presets(
    trailers: Mapping[str, str],
    labels: Sequence[str],
    is_pr: bool,
    is_llvm_integrate_pr: bool,
) -> str:
    """Parses and validates the benchmark presets from trailers.

    Args:
      trailers: trailers from PR description.
      labels: list of PR labels.
      is_pr: is pull request event.
      is_llvm_integrate_pr: is LLVM integration PR.

    Returns:
      A comma separated preset string, which later will be parsed by
      build_tools/benchmarks/export_benchmark_config.py.
    """

    skip_llvm_integrate_benchmark = SKIP_LLVM_INTEGRATE_BENCHMARK_KEY in trailers
    if skip_llvm_integrate_benchmark:
        print(
            "Skipping default benchmarking on LLVM integration because PR "
            f"description has '{SKIP_LLVM_INTEGRATE_BENCHMARK_KEY}' trailer."
        )

    if not is_pr:
        preset_options = {DEFAULT_BENCHMARK_PRESET}
        print(f"Using benchmark presets '{preset_options}' for non-PR run")
    elif is_llvm_integrate_pr and not skip_llvm_integrate_benchmark:
        # Run all benchmark presets for LLVM integration PRs.
        preset_options = {DEFAULT_BENCHMARK_PRESET}
        print(f"Using benchmark preset '{preset_options}' for LLVM integration PR")
    else:
        preset_options = set(
            label.split(":", maxsplit=1)[1]
            for label in labels
            if label.startswith(BENCHMARK_LABEL_PREFIX + ":")
        )
        trailer = trailers.get(BENCHMARK_EXTRA_KEY)
        if trailer is not None:
            preset_options = preset_options.union(
                option.strip() for option in trailer.split(",")
            )
        print(f"Using benchmark preset '{preset_options}' from trailers and labels")

    if DEFAULT_BENCHMARK_PRESET in preset_options:
        preset_options.remove(DEFAULT_BENCHMARK_PRESET)
        preset_options.update(DEFAULT_BENCHMARK_PRESET_GROUP)

    if preset_options.intersection(DEFAULT_BENCHMARK_PRESET_GROUP):
        # The is a sugar to run the compilation benchmarks when any default
        # benchmark preset is present.
        preset_options.add("comp-stats")

    preset_options = sorted(preset_options)
    for preset_option in preset_options:
        if preset_option not in BENCHMARK_PRESET_OPTIONS:
            raise ValueError(
                f"Unknown benchmark preset option: '{preset_option}'.\n"
                f"Available options: '{BENCHMARK_PRESET_OPTIONS}'."
            )

    return ",".join(preset_options)


def main():
    is_pr = os.environ["GITHUB_EVENT_NAME"] == "pull_request"
    trailers, labels = get_trailers_and_labels(is_pr)
    is_llvm_integrate_pr = bool(
        LLVM_INTEGRATE_TITLE_PATTERN.search(os.environ.get("PR_TITLE", ""))
        or LLVM_INTEGRATE_BRANCH_PATTERN.search(os.environ.get("PR_BRANCH", ""))
        or LLVM_INTEGRATE_LABEL in labels
    )
    output = {
        "should-run": json.dumps(should_run_ci(is_pr, trailers)),
        "is-pr": json.dumps(is_pr),
        "runner-env": get_runner_env(trailers),
        "runner-group": "presubmit" if is_pr else "postsubmit",
        "write-caches": "0" if is_pr else "1",
        "benchmark-presets": get_benchmark_presets(
            trailers, labels, is_pr, is_llvm_integrate_pr
        ),
    }

    set_output(output)


if __name__ == "__main__":
    main()