import json

import pandas as pd
import requests
from loguru import logger


class DataPuller(object):
    def __init__(self, db_host, column_rename, payload, absolute_time=None, optional_payload=None):
        self.optional_payload = optional_payload
        self.db_host_url = db_host
        self.request_url = "{kairos_host}/api/v1/datapoints/query".format(kairos_host=self.db_host_url)
        self.payload = payload
        self.column_rename = column_rename
        if absolute_time is not None:
            if "start_relative" in self.payload:
                del self.payload["start_relative"]
            if "end_relative" in self.payload:
                del self.payload["end_relative"]
            self.payload["start_absolute"] = absolute_time["start_absolute"]
            self.payload["end_absolute"] = absolute_time["end_absolute"]

    def get_data(self, start_timestamp, end_timestamp):
        logger.info("Data for the parameters being pulled from Kairos Database")
        self.payload['start_absolute'] = start_timestamp
        self.payload['end_absolute'] = end_timestamp
        response_data = requests.post(url=self.request_url, data=json.dumps(self.payload)).json()
        output_data = response_data["queries"]
        logger.debug("Data pull complete")
        df_final = pd.DataFrame()
        for i in range(len(output_data)):
            grouped_output_data = output_data[i]["results"]
            for each_grouped_data in grouped_output_data:
                value = (each_grouped_data["values"])
                tag_id = each_grouped_data["group_by"][0]["group"]["c3"]
                try:
                    logger.debug(f"Renamed {tag_id} to {self.column_rename[tag_id]} in Data")
                    column_name = self.column_rename[tag_id]
                except KeyError as ke:
                    logger.debug(f"Column Renaming Logic not found for {tag_id} - {ke}")
                    column_name = tag_id
                df_column_data = pd.DataFrame(data=value, columns=["timestamp", column_name])
                if df_final.empty:
                    df_final = df_column_data
                else:
                    df_final = df_final.merge(df_column_data, how="outer", left_on="timestamp", right_on="timestamp")
        df_final["epoch_time"] = df_final["timestamp"]
        df_final["timestamp"] = pd.to_datetime(df_final['timestamp'], unit="ms").dt.tz_localize('UTC').dt.tz_convert(
            'Asia/Kolkata')
        logger.debug(f"Final number of columns : {len(list(df_final.columns))}")
        return df_final
