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

import argparse
import functools
import re
import sys
from argh import arg  # type: ignore
from os import chdir
from pathlib import Path
from typing import List, Optional, Tuple

from impl.common import CROSVM_ROOT, GerritChange, cmd, confirm, run_commands

USAGE = """\
./tools/cl [upload|rebase|status|prune]

Upload changes to the upstream crosvm gerrit.

Multiple projects have their own downstream repository of crosvm and tooling
to upload to those.

This tool allows developers to send commits to the upstream gerrit review site
of crosvm and helps rebase changes if needed.

You need to be on a local branch tracking a remote one. `repo start` does this
for AOSP and chromiumos, or you can do this yourself:

    $ git checkout -b mybranch --track origin/main

Then to upload commits you have made:

    [mybranch] $ ./tools/cl upload

If you are tracking a different branch (e.g. aosp/main or cros/chromeos), the upload may
fail if your commits do not apply cleanly. This tool can help rebase the changes, it will
create a new branch tracking origin/main and cherry-picks your commits.

    [mybranch] $ ./tools/cl rebase
    [mybranch-upstream] ... resolve conflicts
    [mybranch-upstream] $ git add .
    [mybranch-upstream] $ git cherry-pick --continue
    [mybranch-upstream] $ ./tools/cl upload

"""

GERRIT_URL = "https://chromium-review.googlesource.com"
CROSVM_URL = "https://chromium.googlesource.com/crosvm/crosvm"
CROSVM_SSO = "sso://chromium/crosvm/crosvm"

git = cmd("git")
curl = cmd("curl --silent --fail")
chmod = cmd("chmod")


class LocalChange(object):
    sha: str
    title: str
    branch: str

    def __init__(self, sha: str, title: str):
        self.sha = sha
        self.title = title

    @classmethod
    def list_changes(cls, branch: str):
        upstream = get_upstream(branch)
        for line in git(f'log "--format=%H %s" --first-parent {upstream}..{branch}').lines():
            sha_title = line.split(" ", 1)
            yield cls(sha_title[0], sha_title[1])

    @functools.cached_property
    def change_id(self):
        msg = git("log -1 --format=email", self.sha).stdout()
        match = re.search("^Change-Id: (I[a-f0-9]+)", msg, re.MULTILINE)
        if not match:
            return None
        return match.group(1)

    @functools.cached_property
    def gerrit(self):
        if not self.change_id:
            return None
        results = GerritChange.query("project:crosvm/crosvm", self.change_id)
        if len(results) > 1:
            raise Exception(f"Multiple gerrit changes found for commit {self.sha}: {self.title}.")
        return results[0] if results else None

    @property
    def status(self):
        if not self.gerrit:
            return "NOT_UPLOADED"
        else:
            return self.gerrit.status


def get_upstream(branch: str = ""):
    try:
        return git(f"rev-parse --abbrev-ref --symbolic-full-name {branch}@{{u}}").stdout()
    except:
        return None


def list_local_branches():
    return git("for-each-ref --format=%(refname:short) refs/heads").lines()


def get_active_upstream():
    upstream = get_upstream()
    if not upstream:
        default_upstream = "origin/main"
        if confirm(f"You are not tracking an upstream branch. Set upstream to {default_upstream}?"):
            git(f"branch --set-upstream-to {default_upstream}").fg()
            upstream = get_upstream()
    if not upstream:
        raise Exception("You are not tracking an upstream branch.")
    parts = upstream.split("/")
    if len(parts) != 2:
        raise Exception(f"Your upstream branch '{upstream}' is not remote.")
    return (parts[0], parts[1])


def prerequisites():
    if not git("remote get-url origin").success():
        print("Setting up origin")
        git("remote add origin", CROSVM_URL).fg()
    if git("remote get-url origin").stdout() not in [CROSVM_URL, CROSVM_SSO]:
        print("Your remote 'origin' does not point to the main crosvm repository.")
        if confirm(f"Do you want to fix it?"):
            git("remote set-url origin", CROSVM_URL).fg()
        else:
            sys.exit(1)

    # Install gerrit commit hook
    hooks_dir = Path(git("rev-parse --git-path hooks").stdout())
    hook_path = hooks_dir / "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 print_branch_summary(branch: str):
    upstream = get_upstream(branch)
    if not upstream:
        print("Branch", branch, "is not tracking an upstream branch")
        print()
        return
    print("Branch", branch, "tracking", upstream)
    changes = [*LocalChange.list_changes(branch)]
    for change in changes:
        if change.gerrit:
            print(" ", change.status, change.title, f"({change.gerrit.short_url()})")
        else:
            print(" ", change.status, change.title)

    if not changes:
        print("  No changes")
    print()


def status():
    """
    Lists all branches and their local commits.
    """
    for branch in list_local_branches():
        print_branch_summary(branch)


def prune(force: bool = False):
    """
    Deletes branches with changes that have been submitted or abandoned
    """
    current_branch = git("branch --show-current").stdout()
    branches_to_delete = [
        branch
        for branch in list_local_branches()
        if branch != current_branch
        and get_upstream(branch) is not None
        and all(
            change.status in ["ABANDONED", "MERGED"] for change in LocalChange.list_changes(branch)
        )
    ]
    if not branches_to_delete:
        print("No obsolete branches to delete.")
        return

    print("Obsolete branches:")
    print()
    for branch in branches_to_delete:
        print_branch_summary(branch)

    if force or confirm("Do you want to delete the above branches?"):
        git("branch", "-D", *branches_to_delete).fg()


def rebase():
    """
    Rebases changes from the current branch onto origin/main.

    Will create a new branch called 'current-branch'-upstream tracking origin/main. Changes from
    the current branch will then be rebased into the -upstream branch.
    """
    branch_name = git("branch --show-current").stdout()
    upstream_branch_name = branch_name + "-upstream"

    if git("rev-parse", upstream_branch_name).success():
        print(f"Overwriting existing branch {upstream_branch_name}")
        git("log -n1", upstream_branch_name).fg()

    git("fetch -q origin main").fg()
    git("checkout -B", upstream_branch_name, "origin/main").fg()

    print(f"Cherry-picking changes from {branch_name}")
    git(f"cherry-pick {branch_name}@{{u}}..{branch_name}").fg()


def upload(
    dry_run: bool = False,
    reviewer: Optional[str] = None,
    auto_submit: bool = False,
    submit: bool = False,
    try_: bool = False,
):
    """
    Uploads changes to the crosvm main branch.
    """
    remote, branch = get_active_upstream()
    changes = [*LocalChange.list_changes("HEAD")]
    if not changes:
        print("No changes to upload")
        return

    print("Uploading to origin/main:")
    for change in changes:
        print(" ", change.sha, change.title)
    print()

    if len(changes) > 1:
        if not confirm("Uploading {} changes, continue?".format(len(changes))):
            return

    if (remote, branch) != ("origin", "main"):
        print(f"WARNING! Your changes are based on {remote}/{branch}, not origin/main.")
        print("If gerrit rejects your changes, try `./tools/cl rebase -h`.")
        print()
        if not confirm("Upload anyway?"):
            return
        print()

    extra_args: List[str] = []
    if auto_submit:
        extra_args.append("l=Auto-Submit+1")
        try_ = True
    if try_:
        extra_args.append("l=Commit-Queue+1")
    if submit:
        extra_args.append(f"l=Commit-Queue+2")
    if reviewer:
        extra_args.append(f"r={reviewer}")

    git(f"push origin HEAD:refs/for/main%{','.join(extra_args)}").fg(dry_run=dry_run)


if __name__ == "__main__":
    chdir(CROSVM_ROOT)
    prerequisites()
    run_commands(upload, rebase, status, prune, usage=USAGE)