#!/usr/bin/python
# -*- coding: utf8 -*-

import sys
import os
import json
import time
import multiprocessing
import re
import tempfile
import traceback
import imp
import logging
import datetime
import pygraphviz as pgv
from kazoo.client import KazooClient
import random

sys.path.insert(0, '/opt/check-db-availability')
sys.path.insert(0, 'opt/check-db-availability')
import check_db_utils

current_date = datetime.datetime.now().strftime("%Y%m%d")
LOG_PATH = u'/var/log/yandex/db-repl-status'
MAIN_LOG_FILE = u'/var/log/yandex/db-repl-status/main.log'
ALLDB_CONFIG = '/etc/yandex-direct/alldb-config.json'
MAX_FRESHNESS = 300
ZOOKEEPER_NODE = '/direct/db-repl-status/%s/graph'

ZK_HOSTS = "ppc-zk-1.da.yandex.ru:2181,ppc-zk-2.da.yandex.ru:2181,ppc-zk-3.da.yandex.ru:2181"
ZK_TIMEOUT = 10
ZK_RETRY = {"max_tries": 6, "delay": 0.3, "max_jitter": 1, "backoff": 2, "ignore_expire": False}
zkh = KazooClient(hosts=ZK_HOSTS, timeout=ZK_TIMEOUT, connection_retry=ZK_RETRY, command_retry=ZK_RETRY)


def set_logger(filename):
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    logging.basicConfig(
        format=u'%(levelname)-8s [%(asctime)s] %(message)s',
        level=logging.INFO,
        filename="%s.%s" % (filename, current_date)
    )


def check_instance(args):
    try:
        instance_name, instance_conf, db_config = args
        set_logger(u"%s/%s.log" % (LOG_PATH, instance_name))
    
        pass_file = tempfile.NamedTemporaryFile(delete=True)
        os.chmod(pass_file.name, 0600)

        db_config_master = ""
        if 'db_config_master_node' in instance_conf:
            db_config_master = check_db_utils.get_db_config_data(db_config, instance_conf['db_config_master_node'])
        user = instance_conf['mysql_user']
        password = check_db_utils.read_file(instance_conf['mysql_pass_file']).strip()
        check_db_utils.write_pass_to_file(pass_file, password)
    
        graph_data = []
    
        for replica in instance_conf['replicas']:
            replica_host = replica['host']
            replica_port = replica.get('mysql_port') or instance_conf['mysql_port']

            attrs = {}
            masters = []

            alive = check_db_utils.is_host_alive(replica_host, replica_port, user, pass_file)[0]

            if instance_conf['type'] == 'xtradb':
                if alive:
                    masters = [
                        {
                            'host': host,
                            'port': port
                        } for host, port in check_db_utils.get_xtradb_slaves(
                            replica_host, replica_port, user, pass_file
                        ) if host != replica_host
                    ]

                    cluster_status, local_state_comment = check_db_utils.get_xtradb_props(
                        replica_host, replica_port, user, pass_file,
                        ['wsrep_cluster_status', 'wsrep_local_state_comment']
                    )   

                    if cluster_status == 'Primary' and local_state_comment == 'Synced':
                        attrs['color'] = "#00ff00"
                    elif local_state_comment in ["Joining", "Waiting on SST"]:
                        attrs['color'] = "#ff0000"
                    elif local_state_comment == "Donor":
                        attrs['color'] = "#0000ff"
                    elif local_state_comment == "Joined":
                        attrs['color'] = "#ffff00"
                    else:
                        attrs['color'] = "#53377a"
                else:
                    attrs['style'] = 'filled'
                    attrs['fillcolor'] = "#ff0000"

                attrs['label'] = u"<%s:%s<BR/><FONT POINT-SIZE='14'>(%s, %s%s)</FONT>>" % (
                    replica_host, replica_port, cluster_status, local_state_comment,
                    u", &#128295;" if replica['maintenance'] else ""
                )

            elif instance_conf['type'] == 'mysql':
                freshness = 0

                if alive:
                    masters = []
                    host, port = check_db_utils.get_master(
                        replica_host, replica_port, user, pass_file
                    )
                    if host:
                        masters.append({'host': host, 'port': port})
                        
                    freshness = check_db_utils.get_freshness_value(
                        instance_conf['project'], instance_conf['type'],
                        replica_host, replica_port, user, pass_file
                    )

                attrs['label'] = u"<%s:%s<BR/><FONT POINT-SIZE='14'>(%s%s)</FONT>>" % (
                    replica_host, replica_port, str(freshness) + "s" if alive else "dead",
                    u", &#128295;" if replica['maintenance'] else ""
                )
                if not alive:
                    attrs['style'] = 'filled'
                    attrs['fillcolor'] = "#ff496c"
                else:
                    attrs['color'] = ("#ff0000"
                                      if freshness > MAX_FRESHNESS
                                      else "#00ff00")

            if db_config_master == replica_host:
                attrs['shape'] = 'doubleoctagon'

            graph_data.append({
                'host': replica_host,
                'port': replica_port,
                'masters': masters,
                'attrs': attrs,
            })
    
        write_graph(instance_name, graph_data, instance_conf['type'])
    except Exception as e:
        write_graph(instance_name, [], instance_conf['type'], error=True)
        logging.exception("unexpected exception")

    return


def write_graph(instance_name, graph_data, db_type, error=False):
    common_graph_props = {
        'strict': False,
        'rankdir': "TB",
        'center': True,
        'fontsize': 17.0,
        'ratio': "compress",
        'fixedsize': True,
        'size': "8"
    }
    juggler_link = "https://juggler.yandex-team.ru/check_details/?host=direct.prod_mysql-health&amp;amp;service=%s" % instance_name
    graph_label = (
        u"<<TABLE BORDER='0'><TR><TD TARGET='_blank' HREF='%s' TITLE='Open in Juggler'>" +
        u"<FONT COLOR='#00008B' FACE='boldfontname' POINT-SIZE='24'><U>%s</U></FONT></TD></TR></TABLE>>"
    ) % (juggler_link, instance_name)

    try:
        if error:
            raise

        graph = pgv.AGraph(
            directed=db_type == "mysql",
            **common_graph_props
        )
    
        graph.node_attr.update(style='bold', margin=0, width=0, height=0)
        graph.edge_attr.update(weight=1000)
    
        for replica in graph_data:
            if replica.get('masters'):
                for master in replica['masters']:
                    graph.add_edge(replica['host'], master['host'], instance_name, label=master['port'])
            else:
                graph.add_node(replica['host'])
    
        for replica in graph_data:
            node = graph.get_node(replica['host'])
            for attr, value in replica['attrs'].items():
                node.attr[attr] = value
    
        graph.add_subgraph(
            graph.nodes(),
            name='cluster',
            label=graph_label
        )
        graph.layout(prog='dot')
    except:
        graph = pgv.AGraph(
            labelloc="t",
            label=graph_label,
            **common_graph_props
        )
        graph.add_node('Exception occured during computing schema, check logs on ppcback', shape='plaintext')
        graph.layout(prog='dot')

    try:
        zkh.start(ZK_TIMEOUT*3)
        zkh.ensure_path(ZOOKEEPER_NODE % instance_name)
        zkh.set(ZOOKEEPER_NODE % instance_name, graph.string().encode('utf-8'))
    except:
        logging.error("FAIL: can't send graph data for %s" % (ZOOKEEPER_NODE % instance_name))
    else:
        logging.info("SUCCESS: graph data was sent")
    finally:
        zkh.stop()


def run():
    set_logger(MAIN_LOG_FILE)
    with open(ALLDB_CONFIG, 'r') as fd:
        config = json.load(fd)

    instances = [
        (instance, config['instances'][instance], check_db_utils.get_db_config(config['instances'][instance], instance))
        for instance in config['instances'] if config['instances'][instance]["type"] in ["mysql", "xtradb"]
    ]

    processes_pool = multiprocessing.Pool(processes=min(len(instances), multiprocessing.cpu_count()))
    processes_pool.map(check_instance, instances)
    processes_pool.close()
    processes_pool.join()


if __name__ == '__main__':
    run()

