from telemetry import *
import gzip
import json
import os
import pandas as pd
import time

# This is an example of using the DataParseLib
if __name__ == "__main__":
    # csv_data = []
    # June_data_path = "./test/test_data"
    June_data_path = "/Users/yinweiling/Downloads/04-06/06"
    if os.path.exists('Report_time-id-slow_query_bucket.csv'):
        os.remove('Report_time-id-slow_query_bucket.csv')
    for root, dirs, files in os.walk(June_data_path, topdown=False):
        for file_name in files:
            # print(file_name)
            if not file_name.endswith('.gz'):
                continue
            data_file = os.path.join(root, file_name)
            # print(data_file)
            with gzip.open(data_file) as fp:
                raw_data_lines = fp.readlines()
                for raw_data in raw_data_lines:
                    try:
                        request = json2pb(json.loads(raw_data))
                        print("RwaData matches the current schema ^_^")
                    except Exception as e:
                        print(e)
                        print("RwaData doesn't match the current schema !")
                    Telemetry_data = DataParse(raw_data)

                    sql_usage_total = Telemetry_data.get_sql_usage_total()
                    slow_query_stats = Telemetry_data.get_slow_query_stats()
                    if not slow_query_stats:
                        continue
                    slow_query_bucket = slow_query_stats['slowQueryBucket']
                    slow_query_key = list(slow_query_bucket.keys())
                    report_timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(Telemetry_data.get_report_timestamp()))
                    tracking_id = Telemetry_data.get_tracking_id()

                    record_line = [report_timestamp, tracking_id]
                    # temp_key = []
                    for k, v in slow_query_bucket.items():
                        # temp_key.append(k)
                        record_line.append(v)
                    # if slow_query_key != temp_key:
                    #     print("problem !!!")
                    # csv_data.append(record_line)
                    if tracking_id == 'bd113e90-4619-4bb8-9e98-da8b1915f82e':
                        print(data_file, '+', tracking_id)
                        import pdb;pdb.set_trace()

                    csv_data = [record_line]
                    record_col_name = ['Report_time', 'id']+slow_query_key
                    sql_usage_record = pd.DataFrame(csv_data, columns=record_col_name)

                    # try:
                    #     sql_usage_record =pd.DataFrame(csv_data, columns=record_col_name)
                    # except Exception as e:
                    #     print(record_col_name)
                    #     print(record_line)
                    #     print(e)
                    #     import pdb;pdb.set_trace()

                    if os.path.exists('Report_time-id-slow_query_bucket.csv'):
                        sql_usage_record.to_csv('Report_time-id-slow_query_bucket.csv', index=False, mode='a',header=False)
                    else:
                        sql_usage_record.to_csv('Report_time-id-slow_query_bucket.csv', index=False)
                        print(file_name)
                    '''
                    if not sql_usage_total:
                        continue
                    report_timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(Telemetry_data.get_report_timestamp()))
                    tracking_id = Telemetry_data.get_tracking_id()
                    # if tracking_id == "66985d2c-43e3-40ee-8f0d-edbfff437040":
                    #     print(data_file)
                    #     import pdb; pdb.set_trace()
                    record_line = [report_timestamp, tracking_id, sql_usage_total]
                    csv_data.append(record_line)
                    # print(record_line)
                    '''
                    print("*"*20)
    # sql_usage_record =pd.DataFrame(csv_data, columns=['Report_time', 'id', 'sql_total'])
    # sql_usage_record.to_csv('Report_time-id-sql_total.csv', index=False)
    print("done!!")