#!/usr/bin/env python3

"""Helper to automatically rebase onto latest commit possible."""

import argparse
import subprocess
import sys


def git(args, **kwargs):
    """Run git."""
    kwargs.setdefault('check', True)
    kwargs.setdefault('capture_output', True)
    kwargs.setdefault('encoding', 'utf-8')
    return subprocess.run(['git'] + args, **kwargs)


def rebase(target):
    """Try to rebase onto |target|."""
    try:
        git(['rebase', target])
        return True
    except KeyboardInterrupt:
        git(['rebase', '--abort'])
        print('aborted')
        sys.exit(1)
    except:
        git(['rebase', '--abort'])
        return False


def rebase_bisect(lbranch, rbranch, behind):
    """Try to rebase branch as close to |rbranch| as possible."""
    def attempt(pos):
        target = f'{rbranch}~{pos}'
        print(f'Rebasing onto {target} ', end='', flush=True)
        print('.', end='', flush=True)
#        git(['checkout', '-f', target])
        print('.', end='', flush=True)
#        git(['checkout', '-f', lbranch])
        print('. ', end='', flush=True)
        ret = rebase(target)
        print('OK' if ret else 'failed')
        return ret

    #"min" is the latest branch commit while "max" is where we're now.
    min = 0
    max = behind
    old_mid = None
    while True:
        mid = min + (max - min) // 2
        if mid == old_mid or mid < min or mid >= max:
            break
        if attempt(mid):
            max = mid
        else:
            min = mid
        old_mid = mid
    print('Done')


def get_ahead_behind(lbranch, rbranch):
    """Return number of commits |lbranch| is ahead & behind relative to |rbranch|."""
    output = git(['rev-list', '--left-right', '--count', f'{lbranch}...{rbranch}']).stdout
    return [int(x) for x in output.split()]


def get_tracking_branch(branch):
    """Return remote branch that |branch| is tracking."""
    merge = git(['config', '--local', f'branch.{branch}.merge']).stdout.strip()
    if not merge:
        return None

    remote = git(['config', '--local', f'branch.{branch}.remote']).stdout.strip()
    if remote:
        if merge.startswith('refs/heads/'):
            merge = merge[11:]
        return f'{remote}/{merge}'
    else:
        return merge


def get_local_branch():
    """Return the name of the local checked out branch."""
    return git(['branch', '--show-current']).stdout.strip()


def get_parser():
    """Get CLI parser."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        '--skip-initial-rebase-latest', dest='initial_rebase',
        action='store_false', default=True,
        help='skip initial rebase attempt onto the latest branch')
    parser.add_argument(
        'branch', nargs='?',
        help='branch to rebase onto')
    return parser


def main(argv):
    """The main entry point for scripts."""
    parser = get_parser()
    opts = parser.parse_args(argv)

    lbranch = get_local_branch()
    print(f'Local branch resolved to "{lbranch}"')
    if not lbranch:
        print('Unable to resolve local branch', file=sys.stderr)
        return 1

    if opts.branch:
        rbranch = opts.branch
    else:
        rbranch = get_tracking_branch(lbranch)
    print(f'Remote branch resolved to "{rbranch}"')

    ahead, behind = get_ahead_behind(lbranch, rbranch)
    print(f'Branch is {ahead} commits ahead and {behind} commits behind')

    if not behind:
        print('Up-to-date!')
    elif not ahead:
        print('Fast forwarding ...')
        git(['merge'])
    else:
        if opts.initial_rebase:
            print(f'Trying to rebase onto latest {rbranch} ... ', end='', flush=True)
            if rebase(rbranch):
                print('OK!')
                return 0
            print('failed; falling back to bisect')
        rebase_bisect(lbranch, rbranch, behind)


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))