import csv
import os
import sys
from datasdk.sources import metrics, slowlog, event
from datasdk.config import config as sdk_config
from diagnosis.dataloader import config
import pandas as pd
import sqlite3


class DataLoader:
    def __init__(self, loader_cfg):
        self.config = loader_cfg

    def load_data(self):
        print("load...")


class MetricsLoader(DataLoader):
    config: config.MetricsLoadConfig

    def __init__(self, loader_cfg, env: sdk_config.Env, o11y_auth0_client_id, o11y_auth0_client_secret):
        super(MetricsLoader, self).__init__(loader_cfg)
        cfg = sdk_config.Config(
            env, o11y_auth0_client_id, o11y_auth0_client_secret)
        self.loader = metrics.Metrics(cfg)

    def load_data(self):
        rule = [sdk_config.Rule(self.config.rule.name, self.config.rule.expr)]
        ruleCfg = sdk_config.RulesConfig(rule)
        try:
            data = self.loader.getMetrics(self.config.cluster_id, ruleCfg, self.config.start_timestamp,
                                          self.config.end_timestamp, self.config.step)
        except:
            raise Exception("get metrics failed", sys.exc_info()[1])
        else:
            if len(data) > 0:
                self.loader.writeToCSV(data, self.config.out_path)
            else:
                if self.config.rule.name == "tiflash_system_current_metric_StoreSizeUsed":
                    write_tiflash_data(self.config.out_path, self.config.start_timestamp,
                                       self.config.end_timestamp, self.config.step)


def write_tiflash_data(out_path, start_ts, end_ts, step):
    with open(out_path, 'w') as f:
        csvwriter = csv.writer(f, delimiter=',',)
        # write header
        csvwriter.writerow(["tiflash_system_current_metric_StoreSizeUsed{}"])
        end = start_ts
        while end <= end_ts:
            csvwriter.writerow(["0"])
            if end == end_ts:
                break
            end = end + step if end + step < end_ts else end_ts


class SlowLogsLoader(DataLoader):
    config: config.SlowLogLoadConfig

    def __init__(self, loader_cfg, env: sdk_config.Env, o11y_auth0_client_id,
                 o11y_auth0_client_secret):
        super(SlowLogsLoader, self).__init__(loader_cfg)
        cfg = sdk_config.Config(
            env, o11y_auth0_client_id, o11y_auth0_client_secret)
        self.loader = slowlog.SlowLog(cfg)

    def load_data(self):
        try:
            data = self.loader.getSlowLog(self.config.cluster_id,
                                          False,
                                          self.config.start_timestamp,
                                          self.config.end_timestamp)
        except:
            raise Exception("get slowlog failed", sys.exc_info()[1])
        else:
            if len(data) > 0:
                with open(self.config.out_path, 'a+') as f:
                    f.write(data)
            # if len(data) > 0:
            #     print('write ')
            #     self.loader.writeToCSV(data, self.config.out_path)


class EventsLoader(DataLoader):
    config: config.EventLoadConfig

    def __init__(self, loader_cfg, token):
        super(EventsLoader, self).__init__(loader_cfg)
        self.loader = event.Event(token)

    def load_data(self):
        try:
            data = self.loader.get_events(self.config.cluster_id,
                                          self.config.start_timestamp,
                                          self.config.end_timestamp)
        except:
            raise Exception("get event failed", sys.exc_info()[1])
        else:
            if len(data) > 0:
                self.loader.write_to_csv(data, self.config.out_path)


class SqliteLoader(DataLoader):
    config: config.SqliteLoadConfig

    def __init__(self, loader_cfg):
        super(SqliteLoader, self).__init__(loader_cfg)

    def load_data(self):
        sqlite_name = os.path.join(self.config.path, self.config.sqlite_name)
        return convert_tables_to_dataframe(sqlite_name)


def convert_tables_to_dataframe(sqlite_name):
    """
    read sqlite file and convert tables to DataFrame
    """
    try:
        conn = sqlite3.connect(sqlite_name)
    except sqlite3.OperationalError:
        raise Exception("open sqlite file {} failed".format(
            sqlite_name), sys.exc_info()[1])
    cur = conn.cursor()
    # get table names
    cur.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = cur.fetchall()
    result = {}
    for t in tables:
        table_name = t[0]
        if table_name in ['slowlog', 'slowlog_with_type']:
            df = pd.read_sql_query(
                "SELECT * FROM {}".format(table_name), con=conn)
            result[table_name] = df
        else:
            df = pd.read_sql_query(
                "SELECT * FROM {}".format(table_name), index_col="timestamp", con=conn)
            result[table_name] = df
    cur.close()
    conn.close()
    return result