#!/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

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")

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"


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 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.")


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


def create_merge_commits(revision: str, max_size: int = 0, create_dry_run: 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)

    # 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

    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():
        if max(dry_run.get_votes("Commit-Queue")) > 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("I am creating a new dry run.")

    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

    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)
        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)