#!/usr/bin/env python3
# Copyright 2022 The Chromium OS Authors. All rights reserved.
# 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.

import functools
import os
import sys
from datetime import date
from typing import Any, cast

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

import json
import re

from impl.common import CROSVM_ROOT, batched, cmd, quoted, run_commands, very_verbose

git = cmd("git")
git_log = git("log --decorate=no --color=never")
curl = cmd("curl --silent --fail")
chmod = cmd("chmod")

GERRIT_URL = "https://chromium-review.googlesource.com"
REPO_URL = "https://chromium.googlesource.com/chromiumos/platform/crosvm"

# Gerrit tags used to identify bot changes.
MERGE_TAG = "crosvm-merge"
DRY_RUN_TAG = "crosvm-merge-dry-run"

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


def is_gce_instance():
    """
    Returns true if we are running in a GCE instance.

    See http://cloud.google.com/compute/docs/instances/detect-compute-engine
    """
    try:
        metadata = curl("-i http://metadata.google.internal").stdout()
        return "Metadata-Flavor: Google" in metadata
    except:
        return False


def strip_xssi(response: str):
    # See https://gerrit-review.googlesource.com/Documentation/rest-api.html#output
    assert response.startswith(")]}'\n")
    return response[5:]


def gerrit_api_get(path: str):
    response = curl(f"{GERRIT_URL}/{path}").stdout()
    return json.loads(strip_xssi(response))


def gerrit_api_post(path: str, body: Any):
    cookiefile = git("config http.cookiefile").stdout()
    response = curl(
        "--cookie",
        cookiefile,
        "-X POST",
        "-H",
        quoted("Content-Type: application/json"),
        "-d",
        quoted(json.dumps(body)),
        f"{GERRIT_URL}/{path}",
    ).stdout()
    if very_verbose():
        print("Response:", response)
    return json.loads(strip_xssi(response))


class Change(object):
    """
    Class to interact with the gerrit /changes/ API.

    For information on the data format returned by the API, see:
    https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#change-info
    """

    id: str
    _data: Any

    def __init__(self, data: Any):
        self._data = data
        self.id = data["id"]

    @functools.cached_property
    def _details(self) -> Any:
        return gerrit_api_get(f"changes/{self.id}/detail")

    @functools.cached_property
    def _messages(self) -> list[Any]:
        return gerrit_api_get(f"changes/{self.id}/messages")

    def get_votes(self, label_name: str) -> list[int]:
        "Returns the list of votes on `label_name`"
        label_info = self._details.get("labels", {}).get(label_name)
        votes = label_info.get("all", [])
        return [cast(int, v.get("value")) for v in votes]

    def get_messages_by(self, email: str) -> list[str]:
        "Returns all messages posted by the user with the specified `email`."
        return [m["message"] for m in self._messages if m["author"].get("email") == email]

    def review(self, message: str, labels: dict[str, int]):
        "Post review `message` and set the specified review `labels`"
        print("Posting on", self, ":", message, labels)
        gerrit_api_post(
            f"changes/{self.id}/revisions/current/review",
            {"message": message, "labels": labels},
        )

    def abandon(self, message: str):
        print("Abandoning", self, ":", message)
        gerrit_api_post(f"changes/{self.id}/abandon", {"message": message})

    @staticmethod
    def query(*queries: str):
        "Returns a list of gerrit changes matching the provided list of queries."
        return [Change(c) for c in gerrit_api_get(f"changes/?q={'+'.join(queries)}")]

    def __str__(self):
        return f"http://crrev.com/c/{self._data['_number']}"

    def pretty_info(self):
        return f"{self} - {self._data['subject']}"


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


def list_active_dry_runs():
    return Change.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("checkout", tracking).fg(quiet=True)
    git("branch -D", branch_name).fg(quiet=True, check=False)
    git("checkout -b", branch_name, "--track", tracking).fg()


def gerrit_prerequisites():
    "Make sure we can upload to gerrit."
    if is_gce_instance():
        # Grab http cookies for accessing GOB. See go/gob-gce
        gcompute_path = f"{os.environ['KOKORO_ARTIFACTS_DIR']}/gcompute-tools"
        git("clone", "https://gerrit.googlesource.com/gcompute-tools", gcompute_path).fg()
        cmd(f"{gcompute_path}/git-cookie-authdaemon", "--no-fork").fg()

        # Setup correct user info for the GCE service account.
        git("config user.name", quoted("Crosvm Bot")).fg()
        git("config user.email", quoted("crosvm-bot@crosvm-packages.iam.gserviceaccount.com")).fg()
    else:
        # Make sure we have http cookies to access gerrit
        cookie_file = git("config http.cookiefile").stdout(check=False)
        if not cookie_file:
            print("Cannot access gerrit without http cookies.")
            print(f"Install one via: {GERRIT_URL}/new-password")
            sys.exit(1)

    # Re-add remote to ensure it is correct. We cannot use the 'origin' provided by Kokoro.
    git("remote remove origin").fg()
    git("remote add origin", REPO_URL).fg()
    git("fetch -q origin").fg()

    # 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):
    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 origin HEAD:refs/for/{target_branch}%{','.join(extra_params)}").fg()
            break
        except:
            continue


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


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

    # Find list of commits to merge, then batch them into smaller merges.
    commits = git_log(f"HEAD..{from_rev}", "--pretty=%H").lines()
    if not commits:
        print("Nothing to merge.")
        return 0

    # Create a merge commit for each batch
    batches = list(batched(commits, max_size)) if max_size > 0 else [commits]
    for i, batch in enumerate(reversed(batches)):
        target = batch[0]
        commit_range = f"{batch[-1]}..{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 main" 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"{REPO_URL}/+log/{commit_range}",
                bug_notes(commit_range),
            ]
        )

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

        # Perfom merge
        git("merge --no-ff", target, "-m", quoted(message), "-m", quoted(trailers)).fg()

    return len(batches)


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 = "origin/main",
    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()

    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 {revision} into origin/{target_branch}")
        setup_tracking_branch("merge-bot-branch", f"origin/{target_branch}")
        if create_merge_commits(revision, max_size, create_dry_run=False) > 0:
            upload_to_gerrit(
                target_branch,
                f"hashtag={MERGE_TAG}",
                "l=Commit-Queue+1",
                *(["l=Bot-Commit+1"] if is_bot else []),
            )


def update_dry_runs(
    revision: str = "origin/main",
    target_branch: str = "chromeos",
    max_size: int = 15,
    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()

    # 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:
            dry_run.review(
                "I think this dry run was SUCCESSFUL.",
                {
                    "Verified": 1,
                    "Bot-Commit": 0,
                },
            )
        elif "Failed builds" 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 {revision} into origin/{target_branch}")
    setup_tracking_branch("merge-bot-branch", f"origin/{target_branch}")
    if create_merge_commits(revision, max_size, create_dry_run=True) > 0:
        upload_to_gerrit(
            target_branch,
            f"hashtag={DRY_RUN_TAG}",
            "l=Commit-Queue+1",
            *(["l=Bot-Commit+1"] if is_bot else []),
        )


run_commands(create_merge_commits, status, update_merges, update_dry_runs)