#!/usr/bin/env python
import os
import sys
import re
import argparse
import subprocess


def get_path(*path, home=False):
    script_dir = os.path.expanduser('~') if home else os.path.abspath('.')
    return os.path.join(script_dir, *path)


def handle_file(name, append_args=None):
    with open(get_path('.ssh', 'config', home=True)) as f:
        data = f.readlines()

    with open(get_path('.ssh', 'config', home=True), 'w') as f:
        flag = True
        for line in data:
            if re.search(r'^Host {}'.format(name), line):
                flag = False
            elif not line.strip():
                flag = True

            if flag:
                f.write(line)

        if append_args is not None:
            f.write(ADD_TEMPLATE.format(**append_args))


def red_color(s):
    return '\033[0;31m' + s + '\033[0m'


def cyan_color(s):
    return '\033[1;36m' + s + '\033[0m'


parser = argparse.ArgumentParser()
parser.add_argument('opt', help='list add mod del cb(clipboard)')
args = parser.parse_args()

op = args.opt

ADD_TEMPLATE = '''Host {name}
    HostName        {host}
    User            {user}
    IdentityFile    {pem}
    Port            {port}
'''

if op == 'list':
    subprocess.Popen("sed -n 'p' ~/.ssh/config | less -p '^Host .*$'", shell=True).communicate()

elif op == 'del':
    name = input('Delete Host: ').strip()
    handle_file(name)
    print('Delete {} successfully.'.format(name))

elif op == 'cb':
    cb = subprocess.check_output(['pbpaste']).decode().strip()

    try:
        pem = re.search(r'''-i\s*(['"])(.*?)\1''', cb)
        port = re.search(r'''-P\s*(\d+?)''', cb)
        connect_str = re.search(r'\b(\w+)(?::.*)?@(.*)\b', cb)
        host = connect_str.group(2)

        connect_args = {
            'name': input('Host[{}]: '.format(connect_str.group(2))) or host,
            'host': host,
            'port': port.group(1) if port else 22,
            'user': connect_str.group(1),
            'pem': get_path(pem.group(2)) if pem else None,
        }
    except AttributeError:
        print('{}: {}'.format(red_color('Check your CLIPBOARD'), cyan_color(cb)), file=sys.stderr)
        sys.exit(0)

    handle_file(connect_args['name'], connect_args)
    print('Append {} successfully.'.format(connect_args['name']))

else:

    connect_args = {
        'name': input('Name: '),
        'host': input('HostName(IP): '),
        'port': input('Port[22]: ') or 22,
        'user': input('User: '),
        'pem': input('Pem: '),
    }

    if op == 'add':
        with open(get_path('.ssh', 'config', home=True), 'a') as f:
            f.write(ADD_TEMPLATE.format(**connect_args))
        print('Append {} successfully.'.format(connect_args['name']))

    elif op == 'mod':
        handle_file(connect_args['name'], connect_args)
        print('Modify {} successfully.'.format(connect_args['name']))

    else:
        raise NotImplementedError
