#!/usr/bin/python
# -*- coding: utf-8 -*-
# vim: set expandtab:tabstop=4:softtabstop=4:shiftwidth=4
# $Id$

import argparse
from kazoo.client import KazooClient
import logging
import pipes
import pwd
import re
import os
import subprocess
import sys
import yaml

def cmd_str(cmd):
    return ' '.join([pipes.quote(s) for s in cmd])

def run():
    desc = u'''

    # на limtest1 залочить java-intapi на версию 1.2345-1
    limtest-up --lock --limtest 1 java-intapi 1.2345-1
    # разлочить java-intapi на limtest1
    limtest-up --unlock --limtest 1 java-intapi
    # обновить java-intapi на limtest1 на версию 1.2345-1, если он не залочен
    limtest-up --limtest 1 java-intapi 1.2345-1
    # обновить java-intapi на всех незалоченных лимтестах на версию 1.2345-1
    limtest-up java-intapi 1.2345-1

    # есть --dry-run
    limtest-up java-intapi 1.2345-1 --dry-run


    ### Продвинутые опции
    # можно переопределять, с какими хостами или какой родительской нодой работать
    limtest-up java-intapi 1.2345-1 --zk-hosts ppcback01e.yandex.ru:2181,ppcback01f.yandex.ru:2181,ppcback01i.yandex.ru:2181
    limtest-up java-intapi 1.2345-1 --zk-root /direct/limtest

    # произвольное имя лимтеста вместо limtestN
    limtest-up java-intapi 1.2345-1 --limtest-name limtest1     # то же самое, что --limtest 1, но длиннее
    limtest-up java-intapi 1.2345-1 --limtest-name my-limtest

    # использовать мастер-соединение ssh
    limtest-up --limtest 1 direct 1.2345-1 -S /tmp/my-ssh-master.sock
    limtest-up --limtest 1 direct 1.2345-1 --ssh-master-socket=/tmp/my-ssh-master.sock
    


    '''
    default_zk_hosts = ','.join([
        'ppcback01e.yandex.ru:2181',
        'ppcback01f.yandex.ru:2181',
        'ppcback01i.yandex.ru:2181',
    ])
    parser = argparse.ArgumentParser(description=desc, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('app')
    parser.add_argument('version', nargs='?')
    parser.add_argument('--overview', action='store_true')
    parser.add_argument('--zk-hosts', default=default_zk_hosts, help=u'хосты zookeeper через запятую, могут быть с портом')
    parser.add_argument('--zk-root', default='/direct/limtest')
    parser.add_argument('--limtest-name')
    parser.add_argument('--limtest', help=u'сокращение: --limtest=1 аналогичен --limtest-name=limtest1')
    parser.add_argument('--log-filename-prefix', default='limtest-up')
    parser.add_argument('--lock', action='store_true')
    parser.add_argument('--unlock', action='store_true')
    parser.add_argument('--dry-run', action='store_true')
    parser.add_argument('-S', '--ssh-master-socket', help=u'использовать мастер-соединение ssh. Пока не работает с выкладкой java не на ТС')
    args = parser.parse_args()
    if args.limtest and args.limtest_name:
        sys.exit(u'ошибка: --limtest и --limtest-name исключают друг друга')
    elif args.limtest:
        limtest_name = 'limtest' + str(args.limtest)
    else:
        limtest_name = args.limtest_name
    if not args.version and not (args.unlock or args.overview):
        sys.exit(u'ошибка: требуется версия')
    if not limtest_name and (args.lock or args.unlock):
        sys.exit(u'ошибка: требуется имя среды (--limtest-name или --limtest)')
    if args.lock and args.unlock:
        sys.exit(u'ошибка: --lock и --unlock исключают друг друга')
    if args.lock and args.overview:
        sys.exit(u'ошибка: --lock и --overview исключают друг друга')
    if args.unlock and args.overview:
        sys.exit(u'ошибка: --unlock и --overview исключают друг друга')

    zk = KazooClient(hosts=args.zk_hosts)
    zk.start()

    app = args.app
    apps_conf = yaml.load(open('/etc/yandex-direct/direct-apps.conf.yaml', 'r').read())
    app_conf = apps_conf['apps'][app]
    version = args.version
    zk_root = args.zk_root
    is_ts = zk_root == '/direct/np'
    is_java_app = app_conf['type'] == 'arcadia-java'
    limtests = []
    if limtest_name:
        limtest_node = args.zk_root + '/' + limtest_name
        if not zk.exists(limtest_node):
            sys.exit(u'ошибка: нет zookeeper-ноды %s' % limtest_node)
        limtests = [limtest_name]
    else:
        limtests = zk.get_children(zk_root)
    if not version and zk_root != '/direct/np':
        version = subprocess.check_output(['direct-version-get', app]).rstrip().split('\n')[0]
    total_errors_cnt = 0
    username = pwd.getpwuid(os.getuid()).pw_name
    for limtest_name in limtests:
        limtest_node = zk_root + '/' + limtest_name
        hosts_node = '%s/hosts/%s' % (limtest_node, app)
        version_node = '%s/versions/%s' % (limtest_node, app)

        hosts = []
        if zk.exists(hosts_node) and zk.exists(version_node):
            limtest_version = zk.get(version_node)[0].rstrip().split('\n')[0]
            if not limtest_version or args.lock or args.unlock or args.overview:
                hosts = zk.get(hosts_node)[0].split('\n')
            hosts = [h for h in hosts if h != '']

            if args.overview:
                for h in hosts:
                    display_version = limtest_version or version or ''
                    print '\t'.join([limtest_name, app, h, display_version])
                continue
            if args.lock:
                sys.stderr.write('%s: setting ZK node "%s" to value "%s"\n' % (limtest_name, version_node, version))
                if not args.dry_run:
                    zk.set(version_node, version)
            if args.unlock:
                sys.stderr.write('%s: setting ZK node "%s" to value "%s"\n' % (limtest_name, version_node, ''))
                if not args.dry_run:
                    zk.set(version_node, '')

        if len(hosts) == 0:
            sys.stderr.write('%s: no hosts to update, skipping\n' % limtest_name)
            continue

        extra_args = []
        if args.dry_run:
            extra_args.append('--dry-run')
        errors = {}
        if is_ts and is_java_app:
            # TODO обобщить распараллеливание выкладки и сделать везде
            os.execvp('parallel-test-update', ['parallel-test-update', '--hosts', ','.join(hosts), '--wait', '--', 'direct-java-deploy.pl', '--host', 'localhost', '--force-yes', app, version])
        for h in hosts:
            sys.stderr.write('%s: updating host %s\n' % (limtest_name, h))
            log_file_name = '/tmp/%s.%s.%s.%s.%s.log' % (args.log_filename_prefix, username, limtest_name, app, h)
            logf = open(log_file_name, 'w+')
            try:
                ssh_cmd = ['ssh', '-o', 'StrictHostKeyChecking=no']
                if args.ssh_master_socket:
                    ssh_cmd += ['-S', args.ssh_master_socket]
                ssh_cmd.append(h)
                if app == 'direct':
                    subprocess.check_call(ssh_cmd + ['sudo', 'direct-deploy.py', app ,version] + extra_args, stdout=logf, stderr=subprocess.STDOUT)
                elif app == 'dna':
                    subprocess.check_call(ssh_cmd + ['sudo', 'apt-get', 'update'], stdout=logf, stderr=subprocess.STDOUT)
                    subprocess.check_call(ssh_cmd + ['sudo', 'apt-get', 'install', '--yes', '--force-yes', 'yandex-direct-dna=%s' % version] + extra_args, stdout=logf, stderr=subprocess.STDOUT)
                elif is_java_app:
                    cmd = ['direct-java-deploy.pl', '--host', h, '--force-yes', app, version]
                    # java-ТС обрабатывается выше
                    subprocess.check_call(cmd + extra_args, stdout=logf, stderr=subprocess.STDOUT)
                else:
                    sys.exit("don't know how to update app: " + app)
            except subprocess.CalledProcessError as e:
                errors[h] = {
                    'exit_code': e.returncode,
                    'cmd': e.cmd,
                    'log_file_name': log_file_name,
                }
        errors_cnt = len(errors.keys())
        total_errors_cnt += errors_cnt
        if errors_cnt > 0:
            for h in sorted(errors.keys()):
                sys.stderr.write('%s: deploy failed for host %s\n  command exited with code %s:\n    %s\n  log: %s\n\n' % (limtest_name, h, errors[h]['exit_code'], cmd_str(errors[h]['cmd']), errors[h]['log_file_name']))
        else:
            sys.stderr.write('%s: OK\n\n' % limtest_name)
    sys.exit(total_errors_cnt)
    


if __name__ == '__main__':
    run()

