import logging
import uuid
from contextlib import closing

import yaml
from psycopg2 import OperationalError, connect, sql

from csm_test_utils.rds_backup.generation.base import BaseDB, logging_configuration

__all__ = ['Pg2DB']


class Pg2DB(BaseDB):
    """Postgres database instance"""

    def _execute_sql(self, sql_query):
        res = []
        with closing(connect(**self.connection)) as connection:
            try:
                with connection.cursor() as cursor:
                    cursor.execute(sql_query)
                    if cursor.description is not None:
                        res = cursor.fetchall()
            except OperationalError:
                logging.exception('Exception occurred when try to execute SQL')
            connection.commit()
        return res

    def _create_table(self, schema_name: str, table_name: str, *columns):
        """Create table in database if it not exists"""

        model_query = sql.SQL('create table if not exists {}.{} '
                              '(id bigint generated by default as identity primary key,{} text);')
        create_table_query = model_query.format(
            sql.Identifier(schema_name),
            sql.Identifier(table_name),
            sql.SQL(', ').join(map(sql.Identifier, columns))
        )
        self._execute_sql(create_table_query)
        logging.info('Table was created (if it not existed before)')

    def fill_table_by_random_values(self, schema_name: str, table_name: str,
                                    range_start: int, range_stop: int,
                                    *columns):
        """Generate random values and insert it into tables"""
        model_query = sql.SQL(
            'insert into {}.{} select generate_series({},{}) as id, md5(random()::text) AS {};')
        record_query = model_query.format(
            sql.Identifier(schema_name),
            sql.Identifier(table_name),
            sql.Literal(range_start),
            sql.Literal(range_stop),
            sql.SQL(', ').join(map(sql.Identifier, columns)))
        self._execute_sql(record_query)
        logging.info('Generate random content and insert it into table')

    def get_database_size(self) -> int:
        """Get database size which returned in bytes"""
        model_query = sql.SQL('select pg_database_size({});')
        get_size_query = model_query.format(sql.Literal(self.database))
        return self._execute_sql(get_size_query)[0][0]

    def run_test(self, src_file: str):
        logging_configuration()
        logging.info('Script starts')
        with open(src_file) as data_file:
            data = yaml.safe_load(data_file)
        count = data['record_count']
        range_start = 0
        schema_name = 'public'
        table_name = str(uuid.uuid4())
        self._create_table(schema_name, table_name, 'content')
        while not self.is_database_fulfilled(data['max_size_in_bytes']):
            range_end = range_start + count
            self.fill_table_by_random_values(
                schema_name, table_name, range_start, range_end, 'content')
            range_start = range_end + 1
            logging.info('Commit session')

        logging.info('Script finished')
