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

import json
import logging
import os
import socket
import subprocess
import sys
from multiprocessing.dummy import Pool as ThreadPool

import direct_juggler.juggler as dj

SCRIPT_NAME = os.path.basename(__file__)
SERVICE_NAME_COMMON = 'scripts.ptkill-check.working'
SERVICE_NAME_SHARD_TEMPLATE = 'scripts.ptkill-check.{0}.ppc{1}'

HOSTNAME = socket.getfqdn()

DBS_SQL_PATH = '/usr/local/bin/dbs-sql'
CHECK_SLEEP_TIME = 61  # выбрано с учётом настроек pt-kill в пакете `yandex-du-ptkill-ppcdata`: 2 * busy-time + 1

ENV = None
DBCONFIG = None


def get_shards():
    logging.debug("dbconfig {0}: {1}".format(ENV, DBCONFIG))
    with open(DBCONFIG) as db_config_fp:
        db_config = json.load(db_config_fp)
        return list(sorted(map(int, db_config['db_config']['CHILDS']['ppc']['CHILDS'].keys())))


def is_ptkill_running_on_shard(shard):
    """
    Проверяет запущен ли pt-kill на указанном шарде.

    Подготавливается специальный запрос с пометкой для pt-kill, отправляется в mysql, если запрос убивается, то
    считаем, что pt-kill запущен на этом шарде.

    Внимание! Функция работает продолжительное время, до `CHECK_SLEEP_TIME` секунд.

    :param shard: шард, работу pt-kill на котором проверям
    :return: True, если pt-kill
    :rtype: bool
    """

    # SQL запрос содержит строку /* pt-kill-me */, которую pt-kill ищет в запросах для прибивания
    # см. пакет `yandex-du-ptkill-ppcdata`. Предполагается, что её можно использовать в комментариях,
    # но для обхода поведения dbs-sql? (mysql-client?), который вырезает комментарии, включаем её в строковой литерал
    if ENV.count("production"):
        command = [DBS_SQL_PATH, 'pr:ppc:{}'.format(shard), '-B',
                   "select SLEEP({}), '/* pt-kill-me */'".format(CHECK_SLEEP_TIME)]
    else:
        command = [DBS_SQL_PATH, 'ts:ppc:{}'.format(shard), '-B',
                   "select SLEEP({}), '/* pt-kill-me */'".format(CHECK_SLEEP_TIME)]

    try:
        output = subprocess.check_output(command)
    except Exception as ex:
        logging.error('Subprocess call error: %s "%s"', ex, ' '.join(command))
        raise
    assert isinstance(output, basestring)

    try:
        # В первой строке заголовки результирующей таблицы, во второй значения "столбцов"
        second_line = output.split("\n")[1]

        # Если `SLEEP` завершится раньше времени, которое передаётся ему аргументом, он выдаст значение `1`, если
        # же прождёт всё указанное время -- `0`. Используем это поведение, чтобы понять был ли запрос убит `pt-kill`ом.
        sleep_result = second_line.split("\t", 1)[0].strip()
        assert sleep_result in ['0', '1']
        return sleep_result == '1'
    except Exception:
        sql_res = output if len(output) < 120 else (output[:100] + "...<truncated>")
        logging.error("Can't parse sql result %r", sql_res)
        raise


def parallel_check_if_ptkill_is_running(shards):
    """
    :return: список результатов по шардам (True|False), в том же порядке, в каком переданы шарды
    """
    pool = ThreadPool(len(shards))
    try:
        result_for_shards = pool.map(is_ptkill_running_on_shard, shards)
        pool.close()
        pool.join()
        return result_for_shards
    except KeyboardInterrupt:
        pool.terminate()
        raise


def make_juggler_event(shard, ptkill_is_running):
    if ptkill_is_running:
        return {
            'service': SERVICE_NAME_SHARD_TEMPLATE.format(ENV, shard),
            'status': 'OK',
            'description': 'OK',
        }
    else:
        return {
            'service': SERVICE_NAME_SHARD_TEMPLATE.format(shard),
            'status': 'CRIT',
            'description': 'pt-kill does not kill query on shard {}'.format(shard),
        }


def run():
    shards = get_shards()
    result_for_shards = parallel_check_if_ptkill_is_running(shards)
    juggler_events = [make_juggler_event(shard, ptkill_is_running)
                      for shard, ptkill_is_running in zip(shards, result_for_shards)]
    dj.queue_events(juggler_events)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')

    # вычисляем в какой среде запущен скрипт
    try:
        envfile = "/etc/yandex/environment.type"
        with open(envfile, 'r') as fd:
            ENV = fd.read().strip()
        if ENV.lower().count("production"):
            DBCONFIG="/etc/yandex-direct/db-config.json"
        else:
            DBCONFIG="/etc/yandex-direct/db-config-np/db-config.test.json"
    except Exception as err:
        logging.error("error read {0}: {1}".format(envfile, err))

    try:
        if ENV is None or DBCONFIG is None:
            logging.critical("unknown ENV/DBCONFIG: {0}/{1}".format(ENV, DBCONFIG))
            sys.exit(1)
        run()
        dj.queue_events([{
            'service': SERVICE_NAME_COMMON,
            'status': 'OK',
            'description': 'OK',
        }])
    except Exception:
        logging.exception("There was an exception")
        sys.exit(1)
