#!/usr/bin/env python

import subprocess
import sys

# Helper wrapper class around common git commands
class Git:
    # Run single git command
    @staticmethod
    def git(cmd, args=[]):
        try:
            return subprocess.check_output(["git", cmd] + args)
        except:
            return None

    @staticmethod
    def branches():
        return [line.strip() for line in Git.git("branch").splitlines()]

    @staticmethod
    def current_branch():
        return next(b[2:] for b in Git.branches() if b.startswith("* "))

def colorprint(text, color, bold):
    attr = []
    if not sys.stdout.isatty():
        print text
        return
    elif color == "red":
        attr.append('31')
    elif color == "cyan":
        attr.append('36')
    else: # Green
        attr.append('32')

    if bold:
        attr.append('1')

    print '\x1b[%sm%s\x1b[0m' % (';'.join(attr), text)

def visit(branch, indent, all_branches, upstreams, current_branch):
    if branch == current_branch:
        colorprint("%s%s *" % (indent, branch), 'cyan', True)
    elif upstreams[branch] == None:
        colorprint("%s%s" % (indent, branch), 'red', True)
    else:
        colorprint("%s%s" % (indent, branch), 'green', False)

    for child in [b for b in all_branches if upstreams[b] == branch]:
        visit(child, indent + "  ", all_branches, upstreams, current_branch)

def main():
    verbose_branches = [line.strip() for line in Git.git("branch", ["-vv"]).splitlines()]

    upstreams = {}

    all_branches = set()

    for verbose_branch in verbose_branches:
        verbose_branch_info = verbose_branch.split()

        name_idx = 1 if verbose_branch_info[0] == "*" else 0

        branch   = verbose_branch_info[name_idx]
        upstream = verbose_branch_info[name_idx+2][1:-1]

        all_branches.add(branch)
        all_branches.add(upstream)

        # Create upstream mappings
        upstreams[branch] = upstream
        if upstream not in upstreams:
            upstreams[upstream] = None

    current_branch = Git.current_branch()

    rootlist = [branch for branch in all_branches if upstreams[branch] == None]
    for root in rootlist:
        visit(root, "", all_branches, upstreams, current_branch)

if __name__ == "__main__":
    main()