#!/usr/bin/env python

#       Licensed to the Apache Software Foundation (ASF) under one
#       or more contributor license agreements.  See the NOTICE file
#       distributed with this work for additional information
#       regarding copyright ownership.  The ASF licenses this file
#       to you under the Apache License, Version 2.0 (the
#       "License"); you may not use this file except in compliance
#       with the License.  You may obtain a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#       Unless required by applicable law or agreed to in writing,
#       software distributed under the License is distributed on an
#       "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#       KIND, either express or implied.  See the License for the
#       specific language governing permissions and limitations
#       under the License.

import re
import shlex
import subprocess
import sys
try:
    import argparse
except ImportError:
    print('requires package argparse or else Python >= 2.7')
    sys.exit()

try:
    from pygments.console import colorize
except ImportError:
    def colorize(color, message):
        return message

usage = "git mr [-a|-r] [--grep PATTERN] [upstream]"


def main():
    parser = argparse.ArgumentParser(prog='git mr', usage=usage)
    parser.add_argument('-r', action='store_true',
                        dest='examine_remote_branches', default=False,
                        help='examine remote branches')
    parser.add_argument('-a', action='store_true',
                        dest='examine_all_branches', default=False,
                        help='examine all branches')
    parser.add_argument('--grep', nargs=1, metavar='PATTERN', default=None,
                        help='only examine branches matching PATTERN')
    parser.add_argument('upstream', nargs='?', default='HEAD',
                        help='the branch to which everything else is compared, defaults to HEAD')
    args = parser.parse_args()

    if args.examine_all_branches:
        merged_branches = git('branch', '-a', '--merged', args.upstream)
        unmerged_branches = git('branch', '-a', '--no-merged', args.upstream)
    elif args.examine_remote_branches:
        merged_branches = git('branch', '-r', '--merged', args.upstream)
        unmerged_branches = git('branch', '-r', '--no-merged', args.upstream)
    else:
        merged_branches = git('branch', '--merged', args.upstream)
        unmerged_branches = git('branch', '--no-merged', args.upstream)

    merged_branches = [line[2:] for line in merged_branches]
    unmerged_branches = [line[2:] for line in unmerged_branches]
    really_unmerged_branches = []

    if args.grep:
        filter = re.compile(args.grep[0])
        merged_branches = [b for b in merged_branches if filter.search(b)]
        unmerged_branches = [b for b in unmerged_branches if filter.search(b)]

    if merged_branches:
        print('Branches contained by %s:' % args.upstream)
        for branch in merged_branches:
            print('  ' + colorize('green', branch))

    if unmerged_branches:
        header_printed = False
        for branch in unmerged_branches:
            commits = ''.join(
                git('cherry', args.upstream, branch, strip_eol=False))
            if commits.find('+') == -1:
                if not header_printed:
                    print('Branches probably merged to %s:' % args.upstream)
                    header_printed = True
                print('  ' + colorize('green', branch))
            else:
                really_unmerged_branches += [branch]

        if really_unmerged_branches:
            print('Branches probably not merged to %s:' % args.upstream)
            for branch in really_unmerged_branches:
                print('  ' + colorize('red', branch))


def git(*args, **kw):
    if len(args) == 1 and isinstance(args[0], basestring):
        argv = shlex.split(args[0])
    else:
        argv = list(args)
    if argv[0] != 'git':
        argv.insert(0, 'git')
    p = subprocess.Popen(argv, stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT)
    rc = p.wait()
    output = p.stdout.readlines()
    if kw.get('strip_eol', True):
        output = [line.rstrip('\n') for line in output]
    return output


if __name__ == '__main__':
    main()