#!/usr/bin/env python3
# Copyright 2014 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Print dependency tree of branches in local repo.

Example:
origin/main
  cool_feature
    dependent_feature
    other_dependent_feature
  other_feature

Branches are colorized as follows:
  * Red - a remote branch (usually the root of all local branches)
  * Cyan - a local branch which is the same as HEAD
    * Note that multiple branches may be Cyan, if they are all on the same
      commit, and you have that commit checked out.
  * Green - a local branch
  * Blue - a 'branch-heads' branch
  * Magenta - a tag
  * Magenta '{NO UPSTREAM}' - If you have local branches which do not track any
    upstream, then you will see this.
"""

import argparse
import collections
import metrics
import subprocess2
import sys

from git_common import (
    current_branch,
    get_branches_info,
    get_config,
    hash_one,
    run,
    tags,
    upstream,
)

import gclient_utils
import git_common
import setup_color
import from_third_party

Fore = from_third_party.import_module("colorama").Fore
Style = from_third_party.import_module("colorama").Style

DEFAULT_SEPARATOR = " " * 4


class OutputManager(object):
    """Manages a number of OutputLines and formats them into aligned columns."""

    def __init__(self):
        self.lines = []
        self.nocolor = False
        self.max_column_lengths = []
        self.num_columns = None

    def append(self, line):
        # All lines must have the same number of columns.
        if not self.num_columns:
            self.num_columns = len(line.columns)
            self.max_column_lengths = [0] * self.num_columns
        assert self.num_columns == len(line.columns)

        if self.nocolor:
            line.colors = [""] * self.num_columns

        self.lines.append(line)

        # Update maximum column lengths.
        for i, col in enumerate(line.columns):
            self.max_column_lengths[i] = max(
                self.max_column_lengths[i], len(col)
            )

    def merge(self, other):
        for line in other.lines:
            self.append(line)

    def as_formatted_string(self):
        return "\n".join(
            l.as_padded_string(self.max_column_lengths)
            for l in self.lines  # noqa: E741
        )


class OutputLine(object):
    """A single line of data.

    This consists of an equal number of columns, colors and separators."""

    def __init__(self):
        self.columns = []
        self.separators = []
        self.colors = []

    def append(self, data, separator=DEFAULT_SEPARATOR, color=Fore.WHITE):
        self.columns.append(data)
        self.separators.append(separator)
        self.colors.append(color)

    def as_padded_string(self, max_column_lengths):
        """ "Returns the data as a string with each column padded to
        |max_column_lengths|."""
        output_string = ""
        for i, (color, data, separator) in enumerate(
            zip(self.colors, self.columns, self.separators)
        ):
            if max_column_lengths[i] == 0:
                continue

            padding = (max_column_lengths[i] - len(data)) * " "
            output_string += color + data + padding + separator

        return output_string.rstrip()


class BranchMapper(object):
    """A class which constructs output representing the tree's branch structure.

    Attributes:
        __branches_info: a map of branches to their BranchesInfo objects which
            consist of the branch hash, upstream, ahead/behind status, frozen
            status, and subject.
        __gone_branches: a set of upstreams which are not fetchable by git
    """

    def __init__(self):
        self.verbosity = 0
        self.maxjobs = 0
        self.show_subject = False
        self.hide_dormant = False
        self.output = OutputManager()
        self.__gone_branches = set()
        self.__branches_info = None
        self.__parent_map = collections.defaultdict(list)
        self.__current_branch = None
        self.__current_hash = None
        self.__tag_set = None
        self.__root = None
        self.__status_info = {}

    def start(self):
        self.__root = git_common.root()
        self.__branches_info = get_branches_info(
            include_tracking_status=self.verbosity >= 1,
            include_subject=self.show_subject or self.verbosity >= 3,
        )
        if self.verbosity >= 2:
            # Avoid heavy import unless necessary.
            from git_cl import get_cl_statuses, color_for_status, Changelist

            change_cls = [
                Changelist(branchref="refs/heads/" + b)
                for b in self.__branches_info.keys()
                if b
            ]
            status_info = get_cl_statuses(
                change_cls,
                fine_grained=self.verbosity > 2,
                max_processes=self.maxjobs,
            )

            # This is a blocking get which waits for the remote CL status to be
            # retrieved.
            for cl, status in status_info:
                self.__status_info[cl.GetBranch()] = (
                    cl.GetIssueURL(short=True),
                    color_for_status(status),
                    status,
                )

        roots = set()

        # A map of parents to a list of their children.
        for branch, branch_info in self.__branches_info.items():
            if not branch_info:
                continue

            if self.__check_cycle(branch):
                continue
            parent = branch_info.upstream
            if not self.__branches_info[parent]:
                # If the parent is not a known branch, it may be an upstream
                # branch like origin/main or it may be gone. Determine which it
                # is, but don't re-query the same parent multiple times.
                if parent not in roots:
                    if not upstream(branch):
                        self.__gone_branches.add(parent)
                    roots.add(parent)

            self.__parent_map[parent].append(branch)

        self.__current_branch = current_branch()
        self.__current_hash = hash_one("HEAD", short=True)
        self.__tag_set = tags()

        if roots:
            for root in sorted(roots):
                self.__append_branch(root, self.output)
        else:
            no_branches = OutputLine()
            no_branches.append("No User Branches")
            self.output.append(no_branches)

    def __check_cycle(self, branch):
        # Maximum length of the cycle is `num_branches`. This limit avoids
        # running into a cycle which does *not* contain `branch`.
        num_branches = len(self.__branches_info)
        cycle = [branch]
        while len(cycle) < num_branches and self.__branches_info[cycle[-1]]:
            parent = self.__branches_info[cycle[-1]].upstream
            cycle.append(parent)
            if parent == branch:
                print(
                    "Warning: Detected cycle in branches: {}".format(
                        " -> ".join(cycle)
                    ),
                    file=sys.stderr,
                )
                return True
        return False

    def __is_invalid_parent(self, parent):
        return not parent or parent in self.__gone_branches

    def __color_for_branch(self, branch, branch_hash):
        if branch == self.__root or branch.startswith("origin/"):
            color = Fore.RED
        elif branch.startswith("branch-heads"):
            color = Fore.BLUE
        elif self.__is_invalid_parent(branch) or branch in self.__tag_set:
            color = Fore.MAGENTA
        elif self.__current_hash.startswith(branch_hash):
            color = Fore.CYAN
        else:
            color = Fore.GREEN

        if branch_hash and self.__current_hash.startswith(branch_hash):
            color += Style.BRIGHT
        else:
            color += Style.NORMAL

        return color

    def __is_dormant_branch(self, branch):
        if "/" in branch:
            return False

        return get_config(f"branch.{branch}.dormant") == "true"

    def __append_branch(self, branch, output, depth=0):
        """Recurses through the tree structure and appends an OutputLine to the
        OutputManager for each branch."""
        child_output = OutputManager()
        for child in sorted(self.__parent_map.pop(branch, ())):
            self.__append_branch(child, child_output, depth=depth + 1)

        is_dormant_branch = (
            self.__is_dormant_branch(branch)
            if (self.hide_dormant or self.verbosity >= 4)
            else False
        )
        if self.hide_dormant and is_dormant_branch and not child_output.lines:
            return

        branch_info = self.__branches_info[branch]
        if branch_info:
            branch_hash = branch_info.hash
        else:
            try:
                branch_hash = hash_one(branch, short=True)
            except subprocess2.CalledProcessError:
                branch_hash = None

        line = OutputLine()

        # The branch name with appropriate indentation.
        suffix = ""
        if branch == self.__current_branch or (
            self.__current_branch == "HEAD" and branch == self.__current_hash
        ):
            suffix = " *"
        branch_string = branch
        if branch in self.__gone_branches:
            branch_string = "{%s:GONE}" % branch
        if not branch:
            branch_string = "{NO_UPSTREAM}"
        main_string = "  " * depth + branch_string + suffix
        line.append(
            main_string, color=self.__color_for_branch(branch, branch_hash)
        )

        # The branch hash.
        if self.verbosity >= 2:
            line.append(branch_hash or "", separator=" ", color=Fore.RED)

        # The branch tracking status.
        if self.verbosity >= 1:
            commits_string = ""
            behind_string = ""
            front_separator = ""
            center_separator = ""
            back_separator = ""
            if branch_info and not self.__is_invalid_parent(
                branch_info.upstream
            ):
                behind = branch_info.behind
                commits = branch_info.commits

                if commits:
                    commits_string = "%d commit" % commits
                    commits_string += "s" if commits > 1 else " "
                if behind:
                    behind_string = "behind %d" % behind

                if commits or behind:
                    front_separator = "["
                    back_separator = "]"

                if commits and behind:
                    center_separator = "|"

            line.append(front_separator, separator=" ")
            line.append(commits_string, separator=" ", color=Fore.MAGENTA)
            line.append(center_separator, separator=" ")
            line.append(behind_string, separator=" ", color=Fore.MAGENTA)
            line.append(back_separator)

        if self.verbosity >= 4:
            line.append(
                " (dormant)" if is_dormant_branch else "          ",
                separator="  ",
                color=Fore.RED,
            )

        # The Rietveld issue associated with the branch.
        if self.verbosity >= 2:
            (url, color, status) = (
                ("", "", "")
                if self.__is_invalid_parent(branch)
                else self.__status_info[branch]
            )
            if self.verbosity > 2:
                line.append(
                    "{} ({})".format(url, status) if url else "", color=color
                )

                is_frozen = branch_info and branch_info.is_frozen
                line.append("[frozen]" if is_frozen else "", color=Fore.CYAN)
            else:
                line.append(url or "", color=color)

        # The subject of the most recent commit on the branch.
        if self.show_subject:
            if not self.__is_invalid_parent(branch):
                if branch_info:
                    line.append(branch_info.subject)
                else:
                    # Fallback for upstream/remote roots not present in local refs/heads.
                    line.append(run("log", "-n1", "--format=%s", branch, "--"))
            else:
                line.append("")

        output.append(line)

        output.merge(child_output)


def print_desc():
    for line in __doc__.splitlines():
        starpos = line.find("* ")
        if starpos == -1 or "-" not in line:
            print(line)
        else:
            _, color, rest = line.split(None, 2)
            outline = line[: starpos + 1]
            outline += (
                getattr(Fore, color.upper()) + " " + color + " " + Fore.RESET
            )
            outline += rest
            print(outline)
    print("")


@metrics.collector.collect_metrics("git map-branches")
def main(argv):
    if gclient_utils.IsEnvCog():
        print(
            "map-branches command is not supported in non-git environment.",
            file=sys.stderr,
        )
        return 1
    setup_color.init()

    if "-h" in argv:
        print_desc()

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-v",
        action="count",
        default=0,
        help=(
            "Pass once to show tracking info, "
            "twice for hash and review url, "
            "thrice for review status, "
            "four times to mark dormant branches"
        ),
    )
    parser.add_argument(
        "--no-color",
        action="store_true",
        dest="nocolor",
        help="Turn off colors.",
    )
    parser.add_argument(
        "-j",
        "--maxjobs",
        action="store",
        type=int,
        help="The number of jobs to use when retrieving review status",
    )
    parser.add_argument(
        "--show-subject",
        action="store_true",
        dest="show_subject",
        help="Show the commit subject.",
    )
    parser.add_argument(
        "--hide-dormant",
        action="store_true",
        dest="hide_dormant",
        help="Hides dormant branches.",
    )

    opts = parser.parse_args(argv)

    mapper = BranchMapper()
    mapper.verbosity = opts.v
    mapper.output.nocolor = opts.nocolor
    mapper.maxjobs = opts.maxjobs
    mapper.show_subject = opts.show_subject
    mapper.hide_dormant = opts.hide_dormant
    mapper.start()
    print(mapper.output.as_formatted_string())
    return 0


if __name__ == "__main__":
    try:
        with metrics.collector.print_notice_and_exit():
            sys.exit(main(sys.argv[1:]))
    except KeyboardInterrupt:
        sys.stderr.write("interrupted\n")
        sys.exit(1)
