import os

import pandas as pd
from sqlalchemy import create_engine

from diagnosis.dataloader.config import SqliteLoadConfig
from diagnosis.dataloader.dataloader import SqliteLoader


def get_sqlite_prefix():
    return prefix_map['sqlite_prefix']


def get_report_prefix():
    return prefix_map['report_prefix']


def get_checkpoint_prefix():
    return prefix_map['checkpoint_prefix']


prefix_map = {
    'sqlite_prefix': 'drhouse/data/sqlite',
    'report_prefix': 'drhouse/data/report',
    'checkpoint_prefix': 'drhouse/checkpoints'
}


def build_sqlite(df: pd.DataFrame, sqlite_name, table_name):
    label_mp = {}
    metrics_mp = {}
    same_label = {}
    for index, row in df.iterrows():
        ts = index
        label = row['labels']
        value = row['value']
        # 收集所有不同的label
        if label not in label_mp:
            if table_name not in ['latency', 'qps', 'latency2']:
                for key in list(label_mp.keys()):
                    if key.lower() == label.lower():
                        same_label[label] = key
            label_mp[label] = True
        # 收集所有的ts
        if ts not in metrics_mp:
            metrics_mp[ts] = {}
        metrics_mp[ts][label] = value

    ts_list = list(metrics_mp.keys())
    ts_list.sort()
    df_mp = {'timestamp': ts_list}

    for k, v in label_mp.items():
        values = []
        for t in ts_list:
            if k in metrics_mp[t]:
                va = metrics_mp[t][k]
                values.append(va)
            else:
                values.append(0)
        df_mp[k] = values
    new_df = pd.DataFrame(df_mp)
    print(new_df.columns)
    for c in new_df.columns:
        if c in list(same_label.keys()):
            new_df[same_label[c]] = new_df[c] + new_df[same_label[c]]
            new_df.drop(columns=[c], inplace=True)
    new_df.set_index('timestamp', inplace=True)
    engine = create_engine('sqlite:///{}'.format(sqlite_name), echo=True)
    new_df.to_sql(name=table_name, con=engine, if_exists='replace')


def build_temp_sqlite(df: pd.DataFrame, temp_sqlite_name, table_name):
    columns = df.columns
    labels = []
    values = []
    ts = []
    for index, row in df.iterrows():
        for c in columns:
            labels.append(c)
            values.append(row[c])
            ts.append(index)
    new_df = pd.DataFrame({'labels': labels, 'value': values, 'timestamp': ts})
    new_df.set_index('labels', inplace=True)
    engine = create_engine('sqlite:///{}'.format(temp_sqlite_name), echo=True)
    new_df.to_sql(name=table_name, con=engine, if_exists='replace')


if __name__ == '__main__':
    path = 'test_path/drhouse/data/sqlite'
    sqlite = '1379661944603414067_metrics_1684972800_1684975500.sqlite'
    sqlite_full_name = os.path.join(path, sqlite)
    loader = SqliteLoader(SqliteLoadConfig('./', sqlite))
    data = loader.load_data()
    print(data)
    df = data['tidb_executor_statement_total']
    # build_sqlite(df, 'test.sqlite', 'latency')
    build_temp_sqlite(df, './test/temp.sqlite', 'tidb_executor_statement_total')
    loader = SqliteLoader(SqliteLoadConfig('./test', 'temp.sqlite'))
    data = loader.load_data()
    print(data)
    df = data['tidb_executor_statement_total']
    build_sqlite(df, './test/test.sqlite', 'tidb_executor_statement_total')