#!/usr/bin/env python3
# Copyright 2022 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This script is used by the CI system to regularly update the merge and dry run changes.
#
# It can be run locally as well, however some permissions are only given to the bot's service
# account (and are enabled with --is-bot).
#
# See `./tools/chromeos/merge_bot -h` for details.
#
# When testing this script locally, use MERGE_BOT_TEST=1 ./tools/chromeos/merge_bot
# to use different tags and prevent emails from being sent or the CQ from being triggered.

from contextlib import contextmanager
import os
from pathlib import Path
import sys
from datetime import date
from typing import List
import random
import string

sys.path.append(os.path.dirname(sys.path[0]))

import re

from impl.common import CROSVM_ROOT, batched, cmd, quoted, run_commands, GerritChange, GERRIT_URL

git = cmd("git")
git_log = git("log --decorate=no --color=never")
curl = cmd("curl --silent --fail")
chmod = cmd("chmod")
dev_container = cmd("tools/dev_container")
mkdir = cmd("mkdir -p")

UPSTREAM_URL = "https://chromium.googlesource.com/crosvm/crosvm"
CROS_URL = "https://chromium.googlesource.com/chromiumos/platform/crosvm"

# Gerrit tags used to identify bot changes.
TESTING = "MERGE_BOT_TEST" in os.environ
if TESTING:
    MERGE_TAG = "testing-crosvm-merge"
    DRY_RUN_TAG = "testing-crosvm-merge-dry-run"
else:
    MERGE_TAG = "crosvm-merge"  # type: ignore
    DRY_RUN_TAG = "crosvm-merge-dry-run"  # type: ignore

# This is the email of the account that posts CQ messages.
LUCI_EMAIL = "chromeos-scoped@luci-project-accounts.iam.gserviceaccount.com"

# Do not create more dry runs than this within a 24h timespan
MAX_DRY_RUNS_PER_DAY = 2


def list_active_merges():
    return GerritChange.query(
        "project:chromiumos/platform/crosvm",
        "branch:chromeos",
        "status:open",
        f"hashtag:{MERGE_TAG}",
    )


def list_active_dry_runs():
    return GerritChange.query(
        "project:chromiumos/platform/crosvm",
        "branch:chromeos",
        "status:open",
        f"hashtag:{DRY_RUN_TAG}",
    )


def list_recent_dry_runs(age: str):
    return GerritChange.query(
        "project:chromiumos/platform/crosvm",
        "branch:chromeos",
        f"-age:{age}",
        f"hashtag:{DRY_RUN_TAG}",
    )


def bug_notes(commit_range: str):
    "Returns a string with all BUG=... lines of the specified commit range."
    return "\n".join(
        set(
            line
            for line in git_log(commit_range, "--pretty=%b").lines()
            if re.match(r"^BUG=", line, re.I) and not re.match(r"^BUG=None", line, re.I)
        )
    )


def setup_tracking_branch(branch_name: str, tracking: str):
    "Create and checkout `branch_name` tracking `tracking`. Overwrites existing branch."
    git("fetch -q cros", tracking).fg()
    git("checkout", f"cros/{tracking}").fg(quiet=True)
    git("branch -D", branch_name).fg(quiet=True, check=False)
    git("checkout -b", branch_name, "--track", f"cros/{tracking}").fg()


@contextmanager
def tracking_branch_context(branch_name: str, tracking: str):
    "Switches to a tracking branch and back after the context is exited."
    # Remember old head. Prefer branch name if available, otherwise revision of detached head.
    old_head = git("symbolic-ref -q --short HEAD").stdout(check=False)
    if not old_head:
        old_head = git("rev-parse HEAD").stdout()
    setup_tracking_branch(branch_name, tracking)
    yield
    git("checkout", old_head).fg()


def gerrit_prerequisites():
    "Make sure we can upload to gerrit."

    # Setup cros remote which we are merging into
    if git("remote get-url cros").fg(check=False) != 0:
        print("Setting up remote: cros")
        git("remote add cros", CROS_URL).fg()
    actual_remote = git("remote get-url cros").stdout()
    if actual_remote != CROS_URL:
        print(f"WARNING: Your remote 'cros' is {actual_remote} and does not match {CROS_URL}")

    # Install gerrit Change-Id hook
    hook_path = CROSVM_ROOT / ".git/hooks/commit-msg"
    if not hook_path.exists():
        hook_path.parent.mkdir(exist_ok=True)
        curl(f"{GERRIT_URL}/tools/hooks/commit-msg").write_to(hook_path)
        chmod("+x", hook_path).fg()


def upload_to_gerrit(target_branch: str, *extra_params: str):
    if not TESTING:
        extra_params = ("r=crosvm-uprev@google.com", *extra_params)
    for i in range(3):
        try:
            print(f"Uploading to gerrit (Attempt {i})")
            git(f"push cros HEAD:refs/for/{target_branch}%{','.join(extra_params)}").fg()
            return
        except:
            continue
    raise Exception("Could not upload changes to gerrit.")


def rename_files_to_random(dir_path: str):
    "Rename all files in a folder to random file names with extension kept"
    print("Renaming all files in " + dir_path)
    file_names = os.listdir(dir_path)
    for file_name in filter(os.path.isfile, map(lambda x: os.path.join(dir_path, x), file_names)):
        file_extension = os.path.splitext(file_name)[1]
        new_name_stem = "".join(
            random.choice(string.ascii_lowercase + string.digits) for _ in range(16)
        )
        new_path = os.path.join(dir_path, new_name_stem + file_extension)
        print(f"Renaming {file_name} to {new_path}")
        os.rename(file_name, new_path)


def create_pgo_profile():
    "Create PGO profile matching HEAD at merge."
    has_kvm = os.path.exists("/dev/kvm")
    if not has_kvm:
        return
    os.chdir(CROSVM_ROOT)
    tmpdirname = "target/pgotmp/" + "".join(
        random.choice(string.ascii_lowercase + string.digits) for _ in range(16)
    )
    mkdir(tmpdirname).fg()
    benchmark_list = list(
        map(
            lambda x: os.path.splitext(x)[0],
            filter(lambda x: x.endswith(".rs"), os.listdir("e2e_tests/benches")),
        )
    )
    print(f"Building instrumented binary, perf data will be saved to {tmpdirname}")
    dev_container(
        "./tools/build_release --build-profile release --profile-generate /workspace/" + tmpdirname
    ).fg()
    print()
    print("List of benchmarks to run:")
    for bench_name in benchmark_list:
        print(bench_name)
    print()
    dev_container("mkdir -p /var/empty").fg()
    for bench_name in benchmark_list:
        print(f"Running bechmark: {bench_name}")
        dev_container(f"./tools/bench {bench_name}").fg()
        # Instrumented binary always give same file name to generated .profraw files, rename to avoid
        # overwriting profile from previous bench suite
        rename_files_to_random(tmpdirname)
    mkdir("profiles").fg()
    dev_container(
        f"cargo profdata -- merge -o /workspace/profiles/benchmarks.profdata /workspace/{tmpdirname}"
    ).fg()
    dev_container("xz -f -9e -T 0 /workspace/profiles/benchmarks.profdata").fg()


####################################################################################################
# The functions below are callable via the command line


def create_merge_commits(
    revision: str, max_size: int = 0, create_dry_run: bool = False, force_pgo: bool = False
):
    "Merges `revision` into HEAD, creating merge commits including at most `max-size` commits."
    os.chdir(CROSVM_ROOT)

    # Find list of commits to merge, then batch them into smaller merges.
    commits = git_log(f"HEAD..{revision}", "--pretty=%H").lines()
    if not commits:
        print("Nothing to merge.")
        return (0, False)
    else:
        commit_authors = git_log(f"HEAD..{revision}", "--pretty=%an").lines()
        if all(map(lambda x: x == "recipe-roller", commit_authors)):
            print("All commits are from recipe roller, don't merge yet")
            return (0, False)

    # Create a merge commit for each batch
    batches = list(batched(commits, max_size)) if max_size > 0 else [commits]
    has_conflicts = False
    for i, batch in enumerate(reversed(batches)):
        target = batch[0]
        previous_rev = git(f"rev-parse {batch[-1]}^").stdout()
        commit_range = f"{previous_rev}..{batch[0]}"

        # Put together a message containing info about what's in the merge.
        batch_str = f"{i + 1}/{len(batches)}" if len(batches) > 1 else ""
        title = "Merge with upstream" if not create_dry_run else f"Merge dry run"
        message = "\n\n".join(
            [
                f"{title} {date.today().isoformat()} {batch_str}",
                git_log(commit_range, "--oneline").stdout(),
                f"{UPSTREAM_URL}/+log/{commit_range}",
                *([bug_notes(commit_range)] if not create_dry_run else []),
            ]
        )

        # git 'trailers' go into a separate paragraph to make sure they are properly separated.
        trailers = "Commit: False" if create_dry_run or TESTING else ""

        # Perfom merge
        code = git("merge --no-ff", target, "-m", quoted(message), "-m", quoted(trailers)).fg(
            check=False
        )
        if code != 0:
            if not Path(".git/MERGE_HEAD").exists():
                raise Exception("git merge failed for a reason other than merge conflicts.")
            print("Merge has conflicts. Creating commit with conflict markers.")
            git("add --update .").fg()
            message = f"(CONFLICT) {message}"
            git("commit", "-m", quoted(message), "-m", quoted(trailers)).fg()
            has_conflicts = True
        # Only uprev PGO profile on Monday to reduce impact on repo size
        # TODO: b/181105093 - Re-evaluate throttling strategy after sometime
        if date.today().weekday() == 0 or force_pgo:
            create_pgo_profile()
            git("add profiles/benchmarks.profdata.xz").fg()
            git("commit --amend --no-edit").fg()

    return (len(batches), has_conflicts)


def status():
    "Shows the current status of pending merge and dry run changes in gerrit."
    print("Active dry runs:")
    for dry_run in list_active_dry_runs():
        print(dry_run.pretty_info())
    print()
    print("Active merges:")
    for merge in list_active_merges():
        print(merge.pretty_info())


def update_merges(
    revision: str,
    target_branch: str = "chromeos",
    max_size: int = 15,
    is_bot: bool = False,
):
    """Uploads a new set of merge commits if the previous batch has been submitted."""
    gerrit_prerequisites()
    parsed_revision = git("rev-parse", revision).stdout()

    active_merges = list_active_merges()
    if active_merges:
        print("Nothing to do. Previous merges are still pending:")
        for merge in active_merges:
            print(merge.pretty_info())
        return
    else:
        print(f"Creating merge of {parsed_revision} into cros/{target_branch}")
        with tracking_branch_context("merge-bot-branch", target_branch):
            count, has_conflicts = create_merge_commits(
                parsed_revision, max_size, create_dry_run=False
            )
            if count > 0:
                labels: List[str] = []
                if not has_conflicts:
                    if not TESTING:
                        labels.append("l=Commit-Queue+1")
                    if is_bot:
                        labels.append("l=Bot-Commit+1")
                upload_to_gerrit(target_branch, f"hashtag={MERGE_TAG}", *labels)


def update_dry_runs(
    revision: str,
    target_branch: str = "chromeos",
    max_size: int = 0,
    is_bot: bool = False,
):
    """
    Maintains dry run changes in gerrit, usually run by the crosvm bot, but can be called by
    developers as well.
    """
    gerrit_prerequisites()
    parsed_revision = git("rev-parse", revision).stdout()

    # Close active dry runs if they are done.
    print("Checking active dry runs")
    for dry_run in list_active_dry_runs():
        cq_votes = dry_run.get_votes("Commit-Queue")
        if not cq_votes or max(cq_votes) > 0:
            print(dry_run, "CQ is still running.")
            continue

        # Check for luci results and add V+-1 votes to make it easier to identify failed dry runs.
        luci_messages = dry_run.get_messages_by(LUCI_EMAIL)
        if not luci_messages:
            print(dry_run, "No luci messages yet.")
            continue

        last_luci_message = luci_messages[-1]
        if "This CL passed the CQ dry run" in last_luci_message or (
            "This CL has passed the run" in last_luci_message
        ):
            dry_run.review(
                "I think this dry run was SUCCESSFUL.",
                {
                    "Verified": 1,
                    "Bot-Commit": 0,
                },
            )
        elif "Failed builds" in last_luci_message or (
            "This CL has failed the run. Reason:" in last_luci_message
        ):
            dry_run.review(
                "I think this dry run FAILED.",
                {
                    "Verified": -1,
                    "Bot-Commit": 0,
                },
            )

        dry_run.abandon("Dry completed.")

    active_dry_runs = list_active_dry_runs()
    if active_dry_runs:
        print("There are active dry runs, not creating a new one.")
        print("Active dry runs:")
        for dry_run in active_dry_runs:
            print(dry_run.pretty_info())
        return

    num_dry_runs = len(list_recent_dry_runs("1d"))
    if num_dry_runs >= MAX_DRY_RUNS_PER_DAY:
        print(f"Already created {num_dry_runs} in the past 24h. Not creating another one.")
        return

    print(f"Creating dry run merge of {parsed_revision} into cros/{target_branch}")
    with tracking_branch_context("merge-bot-branch", target_branch):
        count, has_conflicts = create_merge_commits(
            parsed_revision, max_size, create_dry_run=True, force_pgo=True
        )
        if count > 0 and not has_conflicts:
            upload_to_gerrit(
                target_branch,
                f"hashtag={DRY_RUN_TAG}",
                *(["l=Commit-Queue+1"] if not TESTING else []),
                *(["l=Bot-Commit+1"] if is_bot else []),
            )
        else:
            if has_conflicts:
                print("Not uploading dry-run with conflicts.")
            else:
                print("Nothing to upload.")


run_commands(
    create_merge_commits,
    status,
    update_merges,
    update_dry_runs,
    gerrit_prerequisites,
)
