import json

from .db_collector import DBCollector
from loguru import logger
import pymysql


class TiDBCollector(DBCollector):
    def __init__(self, db_url, user_name, password):
        self.VERSION_SQL = "SELECT @@GLOBAL.version;"

        # sql for collecting parameters
        self.PARAMETERS_SQL = "SHOW CONFIG;"
        self.VARIABLES_SQL = "SHOW GLOBAL VARIABLES;"

        # sql for getting metrics
        self.METRICS_SQL = {
            "tidb_qps": "SELECT instance,sum(value)/7 'value' from METRICS_SCHEMA.tidb_qps where result='OK' "
                        "and time between now()-interval 4 minute and now()-interval 1 minute group by instance;",
            "tidb_query_duration": "SELECT instance,sum(value)/7 'value' from METRICS_SCHEMA.tidb_query_duration"
                                   " where quantile=0.99 and time between now()-interval 4 minute and now()-interval"
                                   " 1 minute group by instance;"
        }

        self.db_parameters = {}
        self.db_metrics = {}
        self.db_component_instances = {}

        database = db_url.split("/")[1]
        host = db_url.split("/")[0].split(":")[0]
        port = int(db_url.split("/")[0].split(":")[1])
        db = pymysql.connect(host=host, user=user_name, password=password,
                             port=port, db=database, charset='utf8')
        cursor = db.cursor()
        try:
            # collect tidb version
            cursor.execute(self.VERSION_SQL)
            out = cursor.fetchall()
            self.version = out[0][0]

            # collect tidb parameters
            cursor.execute(self.PARAMETERS_SQL)
            out = cursor.fetchall()
            for row in out:
                self.db_parameters["(" + ",".join(row[0:3]) + ")"] = str(row[3])
                if row[0] in self.db_component_instances.keys():
                    self.db_component_instances[row[0]].append(row[1])
                else:
                    self.db_component_instances[row[0]] = []
                    self.db_component_instances[row[0]].append(row[1])

            cursor.execute(self.VARIABLES_SQL)
            out = cursor.fetchall()
            for row in out:
                self.db_parameters["(" + row[0] + ")"] = str(row[1])

            # collect tidb internal metrics
            # step 1: set tidb_metric_query_range_duration and tidb_metric_query_step to 30 sec
            cursor.execute("set@@tidb_metric_query_range_duration=30;")
            cursor.execute("set@@tidb_metric_query_step=30;")

            # step 2: get metrics: ((name, instance), value)
            for key, value in self.METRICS_SQL.items():
                cursor.execute(value)
                out = cursor.fetchall()
                for row in out:
                    self.db_metrics["(" + key + "," + row[0] + ")"] = str(row[1])
        except Exception as e:
            logger.error(e)
            db.rollback()
        finally:
            db.close()

    def collect_parameters(self):
        parameters = {"global": {"global": self.db_parameters}, "local": None}
        return json.dumps(parameters)

    def collect_metrics(self):
        metrics = {"global": {"global": self.db_metrics}, "local": None}
        return json.dumps(metrics)

    def collect_version(self):
        return self.version

    def collect_component_instances(self):
        return self.db_component_instances