import psycopg2
import json


class DatabaseHandler:
    def __init__(self, config):
        self.config = config
        self.connection = psycopg2.connect(database=config['db_database'], user=config['db_user'],
                                           password=config['db_password'], host=config['db_host'], port=config['db_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 ensure_connection(self):
        try:
            self.cursor.execute(f"select tablename from pg_tables where schemaname='public'")
        except Exception as e:
            print(e)
            config = self.config
            self.connection = psycopg2.connect(database=config['db_database'], user=config['db_user'],
                                               password=config['db_password'], host=config['db_host'],
                                               port=config['db_port'])
            self.cursor = self.connection.cursor()

    def create_database(self, db_name, force_create=False):
        self.ensure_connection()
        if force_create:
            self.drop_database(db_name)
        try:
            self.connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
            self.cursor.execute(f'CREATE DATABASE {db_name}')
            print(f'Database {db_name} successfully created.')
        except psycopg2.errors.DuplicateDatabase as e:
            print(f'Database {db_name} already exists: {e}')

    def drop_database(self, db_name):
        self.ensure_connection()
        self.connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
        self.cursor.execute(f'drop database if exists {db_name}')
        print(f'Database {db_name} successfully dropped.')
        self.table_list.remove(db_name)

    def create_table(self, table_structure, table_name, force_create=False, primary_key_list=None):
        self.ensure_connection()
        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):
        self.ensure_connection()
        table_name = table_name.lower()
        if not self.exists_table(table_name):
            if info_flag:
                print(f'table: {table_name} not exists in database.')
            return
        self.cursor.execute(f'drop table {table_name}')
        self.connection.commit()
        if info_flag:
            print(f'table: {table_name} dropped')
        return

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

        :param schema:
        :param batch:
        :param table_name:
        :param data_dict_list:
        :return:
        """
        self.ensure_connection()
        full_table_name = f"{schema}.{table_name.lower()}"
        # 检查表是否存在时也要加上schema
        self.cursor.execute(
            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema=%s AND table_name=%s)",
            (schema, table_name.lower())
        )
        if not self.cursor.fetchone()[0]:
            print(f'table: {full_table_name} not exists in database.')
            return

        table_structure = data_dict_list[0]
        content_list = [f'{key.replace(" ", "_").replace("-", "_")}'.lower() for key in table_structure.keys()]
        sql = f"INSERT INTO {full_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()
            print(f'Inserted into table {full_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:
        """
        self.ensure_connection()
        table_name = table_name.lower()
        if not self.exists_table(table_name):
            print(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')
        if unique_columns:
            sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) "
                   f"VALUES({','.join(['%s'] * len(table_structure))})"
                   f" ON CONFLICT ({unique_columns}) DO NOTHING")
        else:
            sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) "
                   f"VALUES({','.join(['%s'] * len(table_structure))})")
        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()
            print(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}')

    def get_data(self, table_name, field_name_list=None):
        self.ensure_connection()
        if not field_name_list:
            sql = f"SELECT * FROM {table_name}"
        elif isinstance(field_name_list, str):
            sql = f"SELECT {field_name_list} FROM {table_name}"
        else:
            sql = f"select {','.join(field_name_list)} from {table_name}"
        self.cursor.execute(sql)
        res = self.cursor.fetchall()
        return res

    def update_table(self, table_name, update_data: dict, key_column):
        self.ensure_connection()
        table_name = table_name.lower()
        if not self.exists_table(table_name):
            print(f'table: {table_name} not exists in database.')
            return
        if key_column not in update_data:
            raise ValueError(f'{key_column} not in data')
        target_list = []
        for key, value in update_data.items():
            if key == key_column:
                continue
            if isinstance(value, dict):
                target_list.append(f"""{key} = '{json.dumps(value).replace("'", "''")}'""")
            else:
                target_list.append(f"""{key} = '{str(value).replace("'", "''")}'""")
        target = ','.join(target_list)
        sql = f"update {table_name} set {target} where {key_column}='{update_data[key_column]}'"
        self.cursor.execute(sql)
        self.connection.commit()

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


if __name__ == "__main__":
    config = {
"db_user": "app",
"db_password": "Letmein1",
"db_host": "pgm-8vb3061igsxn50is9o.pgsql.zhangbei.rds.aliyuncs.com",
"db_port": "5432",
"db_database": "flex_desk",
}
    # db_handler = DatabaseHandler(config)
    # print(db_handler.get_data("wiki.feishu_doc_public"))
    # db_handler.insert_into_table('feishu_doc_public', [table_schema], batch=1, schema='wiki')
