from thefuck.utils import replace_argument
from thefuck.specific.git import git_support


@git_support
def match(command):
    return ('push' in command.script
            and 'set-upstream' in command.stderr)


@git_support
def get_new_command(command):
    # If --set-upstream or -u are passed, remove it and its argument. This is
    # because the remaining arguments are concatenated onto the command suggested
    # by git, which includes --set-upstream and its argument
    upstream_option_index = -1
    try:
        upstream_option_index = command.script_parts.index('--set-upstream')
    except ValueError:
        pass
    try:
        upstream_option_index = command.script_parts.index('-u')
    except ValueError:
        pass
    if upstream_option_index is not -1:
        command.script_parts.pop(upstream_option_index)
        command.script_parts.pop(upstream_option_index)

    push_upstream = command.stderr.split('\n')[-3].strip().partition('git ')[2]
    return replace_argument(" ".join(command.script_parts), 'push', push_upstream)