#!/usr/bin/env python
# *-* encoding: utf-8 *-*

import argparse
import os
import socket
import datetime
import json
import logging
import direct_juggler.juggler as dj
from kazoo.client import KazooClient

SERVICE_NAME = 'zk_cluster_alive.%s'
DESCRIPTION_PREFIX = 'from: %s' % os.path.basename(__file__)

ZK_NODE_PATH = '/heartbeat/%s' % socket.getfqdn()
ZK_TIMEOUT = 5
ZK_RETRY = {"max_tries": 0, "ignore_expire": False}

DB_CONFIG_PATH = '/etc/yandex-direct/db-config.json'
with open(DB_CONFIG_PATH, 'r') as fh:
    DB_CONFIG = json.loads(fh.read())


def parse_options():
    parser = argparse.ArgumentParser(description="Проверяет доступность zookeeper-кластера")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        '-i', '--instance', type=str,
        choices=[ins for ins, val in DB_CONFIG['db_config']['CHILDS'].items() if val.get('engine', '') == 'zookeeper'],
        help="название инстанса (хосты берутся из дб-конфига)"
    )
    group.add_argument(
        '--hosts', type=str,
        help="вручную указать хосты кластера через запятую"
    )
    parser.add_argument(
        '-j', '--juggler', action='store_true',
        help='отправить результат в juggler и ничего не выводить',
    )
    opts, extra = parser.parse_known_args()

    if len(extra) > 0:
        die("There are unknown parameters: %s" % " ".join(extra))

    if opts.juggler and opts.hosts:
        parser.error("нельзя использовать --hosts и --juggler одновременно")

    return opts


def run():
    opts = parse_options()
    if not opts.juggler:
        logging.basicConfig(level=logging.INFO)
    else:
        logging.basicConfig(level=logging.CRITICAL)

    is_crit = False
    if opts.instance:
        hosts = ",".join(DB_CONFIG['db_config']['CHILDS'][opts.instance]['host'])
    else:
        hosts = opts.hosts

    try:
        zkh = KazooClient(hosts=hosts, timeout=ZK_TIMEOUT, connection_retry=ZK_RETRY, command_retry=ZK_RETRY)
        zkh.start(timeout=ZK_TIMEOUT * hosts.count(',') + 1)
        zkh.ensure_path(ZK_NODE_PATH)

        cur_dt = str(datetime.datetime.now())
        zkh.set(ZK_NODE_PATH, cur_dt)
        if cur_dt != zkh.get(ZK_NODE_PATH)[0]:
            is_crit = True

    except Exception as e:
        is_crit = True
        if not opts.juggler:
            logging.exception('exception')

    if opts.juggler:
        dj.queue_events([{
            'service': SERVICE_NAME % opts.instance,
            'status': 'CRIT' if is_crit else 'OK',
            'description': "; ".join([
                DESCRIPTION_PREFIX,
                'run this script for details',
            ]),
        }])
    else:
        print '-' * 40
        print 'CRIT: cluster is down' if is_crit else 'OK'

    return


if __name__ == '__main__':
    run()

