#!/usr/bin/python
# -*- coding: utf-8 -*-

import argparse
import logging
import os
import pipes
import subprocess
import yaml
import sys
import time

def run():
    script_name = os.path.basename(__file__)
    parser = argparse.ArgumentParser(usage='{} [--debug] [--sleep N] -- cmd [arg1] [arg2] ...'.format(script_name),
                                     description=u'Запустить команду, закрыв предварительно приложение от балансера, и открыв после этого.\nЗакрытие от балансера осуществляется созданием специального файла (настройка alive_force_fail_file в apps-конфиге), наличие которого должно обрабатываться в nginx. Если файл уже существовал до запуска этой команды, то он не удаляется в конце (сохраняется состояние закрытости).')
    parser.add_argument('app', help=u'приложение из apps-конфига')
    parser.add_argument('--debug', action='store_true', required=False)
    parser.add_argument('--sleep', metavar='N', type=int, default=10, required=False, help=u'сколько ждать в секундах после создания и удаления файла. Нужно, чтобы балансер "осознал" изменение состояния перед запуском команды. По умолчанию 10 секнуд.')
    # parse_known_args не останавливается на первом неопознанном параметре
    # так что <script_name> cmd --debug запустит не 'cmd --debug', а cmd, и скрипт будет логировать с уровнем debug
    # поэтому usage безусловно рекомендует отделять команду с помощью --
    # может быть, стоит отказаться от argparse
    (args, cmd) = parser.parse_known_args()
    if args.debug:
        logging.basicConfig(format='%(asctime)s %(message)s', level=logging.DEBUG)
    else:
        logging.basicConfig(format='%(message)s', level=logging.INFO)
    cmd_str = ' '.join([pipes.quote(s) for s in cmd])
    app = args.app
    sleep_delay = args.sleep    # секунд
    with open('/etc/yandex-direct/direct-apps.conf.yaml', 'r') as f:
        apps_config = yaml.load(f)
        if app not in apps_config['apps']:
            sys.exit("unknown app '{}', exiting".format(app))
        alive_force_fail_file = apps_config['apps'][app].get('alive_force_fail_file')
    alive_force_fail_file_existed = True
    if alive_force_fail_file:
        if not os.path.exists(alive_force_fail_file):
            alive_force_fail_file_existed = False
            logging.debug("creating file '{}'".format(alive_force_fail_file))
            with open(alive_force_fail_file, 'w+') as f:
                f.write('')
        else:
            logging.debug("file '{}' already exists".format(alive_force_fail_file))
        logging.debug('sleeping for {} seconds'.format(sleep_delay))
        time.sleep(sleep_delay)
    logging.debug("running command: {}".format(cmd_str))
    cmd_exit_code = subprocess.call(cmd)
    # другим вариантом было бы не открывать от балансера при падении, но так как неизвестно, что делает команда, это опасно.
    if cmd_exit_code > 0:
        logging.warn("command '{}' exited with non-zero exit code {}".format(cmd_str, cmd_exit_code))
    if alive_force_fail_file and not alive_force_fail_file_existed:
        logging.debug("removing file '{}'".format(alive_force_fail_file))
        os.remove(alive_force_fail_file)
        logging.debug("sleeping for {} seconds".format(sleep_delay))
        time.sleep(sleep_delay)
    sys.exit(cmd_exit_code)

if __name__ == '__main__':
    run()
