import os
import json
import yaml
import MySQLdb
import geoip2.database
from datetime import datetime
from pyspark import SparkContext, SparkConf

from settings import (
    DB_HOST, DB_PORT, DB_USER, DB_PASSWD, DB_NAME, BASE_DIR, AWS_LOGS_BUCKET, AWS_REGION, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY
)
from utils.spark import SparkTaskAmazonS3


def get_cluster_ip(log):
    log = json.loads(log)
    ip = log['ip']
    cluster_id = log['content']['EventDetail']['Cluster']['UUID']
    return (cluster_id, ip)


def get_cluster_id(log):
    log = json.loads(log)
    return log['content']['EventDetail']['Cluster']['UUID']


def parse_node_num(log):
    log = json.loads(log)
    cluster_id = log['content']['EventDetail']['Cluster']['UUID']
    command = log['content']['EventDetail']['Cluster'].get('command')
    topology = log['content']['EventDetail']['Cluster'].get('topology')
    # only deploy and scale-out is setted
    if not (command and topology):
        return None
    if not command.startswith('cluster deploy'):
        return None
    try:
        # topology is set by user itself, so it may invalid
        topology = yaml.load(topology, Loader=yaml.FullLoader)
    except Exception:
        return None
    pd_servers, pd_count, pd_hosts = topology.get('pd_servers'), 0, []
    tikv_servers, tikv_count, tikv_hosts = topology.get('tikv_servers'), 0, []
    tidb_servers, tidb_count, tidb_hosts = topology.get('tidb_servers'), 0, []
    if pd_servers:
        try:
            pd_hosts = [x['host'] for x in pd_servers]
        except Exception:
            pd_hosts = []
        pd_count = len(set(pd_hosts))
    if tikv_servers:
        try:
            tikv_hosts = [x['host'] for x in tikv_servers]
        except Exception:
            tikv_hosts = []
        tikv_count = len(set(tikv_hosts))
    if tidb_servers:
        try:
            tidb_hosts = [x['host'] for x in tidb_servers]
        except Exception:
            tidb_hosts = []
        tidb_count = len(set(tidb_hosts))
    total_count = len(set(pd_hosts + tikv_hosts + tidb_hosts))
    return (cluster_id, (total_count, pd_count, tikv_count, tidb_count))


def get_max_node_num(nodes_a, nodes_b):
    if nodes_a[0] > nodes_b[0]:
        return nodes_a
    return nodes_b


class TiupSparkTask(SparkTaskAmazonS3):

    def __init__(self, day: str):
        super().__init__(
            app_name='telemetry-tiup-report',
            day=day,
            bucket=AWS_LOGS_BUCKET,
            region=AWS_REGION,
            access_key_id=AWS_ACCESS_KEY_ID,
            secret_access_key=AWS_SECRET_ACCESS_KEY,
            path_pattern='/tiup/%Y/%m/%d/*/*'
        )
        self.target_day = datetime.strptime(day, '%Y-%m-%d')
        self.target_day_str = day

    def run(self):
        logs = self.files()
        rdd_cluster_id_ip_collection = logs.map(get_cluster_ip)
        rdd_cluster_ids = logs.map(get_cluster_id)
        rdd_cluster_event_num_collection = rdd_cluster_ids.map(
            lambda l: (l, 1)).reduceByKey(lambda x, y: x + y)
        rdd_cluster_node_num_collection = logs.map(parse_node_num)\
                                              .filter(lambda x: x)\
                                              .reduceByKey(get_max_node_num)

        cluster_ids = rdd_cluster_ids.distinct().collect()
        cluster_ip_map = dict(rdd_cluster_id_ip_collection.collect())
        cluster_event_num_map = dict(rdd_cluster_event_num_collection.collect())
        cluster_node_num_map = dict(rdd_cluster_node_num_collection.collect())

        db = MySQLdb.connect(
            host=DB_HOST,
            port=DB_PORT,
            user=DB_USER,
            passwd=DB_PASSWD,
            db=DB_NAME,
            ssl={}
        )
        cursor = db.cursor()

        # 取出已有集群 ID
        cursor.execute("""
            SELECT distinct(`cluster_id`) FROM `spark_tiup_clusters`
            """)
        exists_cluster_ids = set(map(lambda x: x[0], cursor.fetchall()))
        new_cluster_ids = set(cluster_ids) - exists_cluster_ids

        # 将新集群入库到 `spark_tiup_clusters`
        new_cluster_data = [
            (i, self.target_day_str, self.target_day_str, '', 0)
            for i in new_cluster_ids
        ]
        cursor.executemany("""
            INSERT INTO
                spark_tiup_clusters (`cluster_id`, `first_login`, `last_login`, `ip`, `event_num`)
            VALUES (%s, %s, %s, %s, %s)
        """, new_cluster_data)
        cursor.execute('commit;')

        # 将今日集群信息入库到 `spark_tiup_cluster_by_date`
        cluster_data = []
        for cluster_id in cluster_ids:
            ip = cluster_ip_map.get(cluster_id)
            event_num = cluster_event_num_map.get(cluster_id)
            cluster_data.append(
                (cluster_id, self.target_day_str, ip, event_num)
            )
        cursor.executemany("""
            INSERT INTO
                spark_tiup_cluster_by_date (`cluster_id`, `date`, `ip`, `event_num`)
            VALUES (%s, %s, %s, %s)
        """, cluster_data)
        cursor.execute('commit;')

        # 更新已有集群的节点数 node_num
        for cluster_id, nodes_count in cluster_node_num_map.items():
            total_count, pd_count, tikv_count, tidb_count = nodes_count
            cursor.execute("""
            UPDATE spark_tiup_clusters
            SET total_node_num = %s, pd_node_num = %s, tikv_node_num = %s, tidb_node_num = %s
            WHERE cluster_id = %s
            """, (total_count, pd_count, tikv_count, tidb_count, cluster_id))
        cursor.execute('commit;')

        geo_reader = geoip2.database.Reader(os.path.join(BASE_DIR, 'static/GeoIP2-City.mmdb'))

        # 更新已有集群的 ip、event_num、last_login 信息
        for d in cluster_data:
            _cluster_id, _last_login, _ip, _event_num = d
            cursor.execute("""
            UPDATE spark_tiup_clusters
            SET ip = %s, geo = %s, event_num = event_num + %s, last_login = %s
            WHERE cluster_id = %s
            """, (
                _ip, json.dumps(geo_reader.city(_ip).raw),
                _event_num, _last_login, _cluster_id))
        cursor.execute('commit;')