import psycopg2
import logging
from src.My_Logger.project_logger import LoggerClass


class DatabaseHandler:
    def __init__(self, config, logger=None):
        if not logger or not isinstance(logger, logging.RootLogger):
            log = LoggerClass()
            self.logger = log.get_log(logging.DEBUG)
        else:
            self.logger = logger
        self.config = config
        # {database="postgres", user="postgres", password="123456", host="localhost", port="5432"}
        self.connection = psycopg2.connect(database=config['database'], user=config['user'],
                                           password=config['password'], host=config['host'], port=config['port'])
        self.cursor = self.connection.cursor()
        self.cursor.execute(f"select tablename from pg_tables where schemaname='public'")
        self.table_list = [item[0] for item in self.cursor.fetchall()]

    def create_table(self, table_structure, table_name, force_create=False, primary_key_list=None):
        table_name = table_name.lower()
        if force_create and self.exists_table(table_name):
            self.drop_table(table_name)
        sql = f'Create Table {table_name}('
        content_list = []
        for key, values in table_structure.items():
            key_string = f'{key.replace(" ", "_").replace("-", "_")}           {values}'
            if primary_key_list and key in primary_key_list:
                key_string += ' PRIMARY KEY'
            content_list.append(key_string)
        sql += ',\n'.join(content_list)
        sql += ');'
        self.connection.cursor().execute(sql)
        self.connection.commit()
        self.table_list.append(table_name)

    def exists_table(self, table_name):
        table_name = table_name.lower()
        if table_name in self.table_list:
            return True
        else:
            return False

    def drop_table(self, table_name, info_flag=False):
        table_name = table_name.lower()
        if not self.exists_table(table_name):
            if info_flag:
                self.logger.error(f'table: {table_name} not exists in database.')
            return
        self.cursor.execute(f'drop table {table_name}')
        self.connection.commit()
        if info_flag:
            self.logger.info(f'table: {table_name} dropped')
        return

    def insert_into_table(self, table_name, data_dict_list, batch=1000):
        """

        :param batch:
        :param table_name:
        :param data_dict_list:
        :return:
        """
        table_name = table_name.lower()
        if not self.exists_table(table_name):
            self.logger.error(f'table: {table_name} not exists in database.')
            return

        table_structure = data_dict_list[0]
        content_list = []
        for key, values in table_structure.items():
            content_list.append(f'{key.replace(" ", "_").replace("-", "_")}'.lower())

        sql = f"INSERT INTO {table_name}({','.join(content_list)}) VALUES({','.join(['%s'] * len(content_list))})"
        target = [list(item.values()) for item in data_dict_list]

        for i in range(0, len(target), batch):
            self.cursor.executemany(sql, target[i: i+batch])
            self.connection.commit()
            self.logger.info(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}')

    def insert_json_into_table(self, table_name, data_list, batch=1000, unique_columns=None):
        """
        insert a list of json into table. The json content should contain all the required columns as its keys.
        :param batch: batch size
        :param table_name:  which table to insert into
        :param data_list:   target json
        :param unique_columns:  the unique column of the table. If set, that column will be
        checked to ensure no duplicate entries
        :return:
        """
        table_name = table_name.lower()
        if not self.exists_table(table_name):
            self.logger.error(f'table: {table_name} not exists in database.')
            return

        table_structure = list(data_list[0].keys())
        if unique_columns and unique_columns not in table_structure:
            raise ValueError(f'{unique_columns} not in data')
        sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) VALUES({','.join(['%s'] * len(table_structure))})"
               f" ON CONFLICT ({unique_columns}) DO NOTHING")
        target = [list(item.values()) for item in data_list]

        for i in range(0, len(target), batch):
            try:
                self.cursor.executemany(sql, target[i: i+batch])
            except psycopg2.errors.InternalError_ as e:
                print(e)
                size = 0
                for line in target[i:i+batch]:
                    if line[-1]:
                        size += line[-1]
                print(size / 1024 / 1024)
                print('error')
            except Exception as e:
                print(sql)
                print(e)
            self.connection.commit()
            self.logger.info(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}')

    def close(self):
        self.connection.close()
