from typing import Any, Dict, List, Optional, Union

import backoff
from psycopg2 import InternalError, OperationalError, sql
from smelog.factory import BoundLogger

from config import POSTGRES_SCHEMA, S3_BUCKET_REGION
from constants import BACKOFF_TIMEOUT, CSV_DELIMITER, CSV_QUOTE_CHAR
from db.pg_connector import PostgresConnector


__all__ = ["PGMixin"]


class PGMixin:
    logger: BoundLogger

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.pg_instance: PostgresConnector = PostgresConnector()
        self._schema_identifier: sql.Identifier = sql.Identifier(POSTGRES_SCHEMA)

    def _pre_run(self):
        super()._pre_run()
        self.pg_instance.init()

    def _post_run(self):
        super()._post_run()
        self.pg_instance.close()

    @backoff.on_exception(backoff.expo, (OperationalError, InternalError), max_time=BACKOFF_TIMEOUT)
    def execute_query(self, query: Union[str, sql.Composable], params: Optional[Dict[str, Any]] = None):
        self.pg_instance.execute_query(query, params)

    @backoff.on_exception(backoff.expo, (OperationalError, InternalError), max_time=BACKOFF_TIMEOUT)
    def execute_select_query(
        self, query: Union[str, sql.Composable], params: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        return self.pg_instance.execute_select_query(query, params)

    @backoff.on_exception(backoff.expo, (OperationalError, InternalError), max_time=BACKOFF_TIMEOUT)
    def swap_tables(self, table: str, table_alt: str):
        self.logger.info(f"Swapping tables '{table_alt}' -> '{table}'")
        with self.pg_instance.transaction() as cursor:
            self.logger.debug(f"Dropping table '{table}'")
            cursor.execute(
                sql.SQL(
                    """
                DROP TABLE IF EXISTS {schema}.{table}
            """
                ).format(schema=self._schema_identifier, table=sql.Identifier(table))
            )

            self.logger.debug(f"Renaming '{table_alt}' to '{table}'")
            cursor.execute(
                sql.SQL(
                    """
                ALTER TABLE {schema}.{table_alt}
                RENAME TO {table}
            """
                ).format(
                    schema=self._schema_identifier, table=sql.Identifier(table), table_alt=sql.Identifier(table_alt)
                )
            )

    def drop_table(self, table: str):
        self.logger.info(f"Dropping table '{table}'")
        self.execute_query(
            sql.SQL(
                """
            DROP TABLE IF EXISTS {schema}.{table}
        """
            ).format(schema=self._schema_identifier, table=sql.Identifier(table))
        )

    def create_table(self, table: str, table_definition: Dict[str, str]):
        self.logger.info(f"Creating '{table}'")

        columns = []
        for key, value in table_definition.items():
            columns.append(sql.SQL("{} {}").format(sql.Identifier(key), sql.SQL(value)))
        columns = sql.SQL(", ").join(columns)

        self.execute_query(
            sql.SQL(
                """
            CREATE TABLE IF NOT EXISTS {schema}.{table}
            (
                {columns}
            )
        """
            ).format(schema=self._schema_identifier, table=sql.Identifier(table), columns=columns)
        )

    def wipe_table(self, table: str):
        self.logger.info(f"Wiping '{table}'")
        self.execute_query(
            sql.SQL(
                """
            DELETE
            FROM {schema}.{table}
        """
            ).format(schema=self._schema_identifier, table=sql.Identifier(table))
        )

    def import_from_s3(self, table: str, bucket_name: str, key: str, **kwargs) -> List[Dict[str, Any]]:
        kwargs.setdefault("delimiter", CSV_DELIMITER)
        kwargs.setdefault("quote_char", CSV_QUOTE_CHAR)
        kwargs.setdefault("bucket_region", S3_BUCKET_REGION)

        table = f"{POSTGRES_SCHEMA}.{table}"

        query = """
            SELECT aws_s3.table_import_from_s3(
               %(table)s,
               '',
               '(format CSV, header True, delimiter E'%(delimiter)s', quote '%(quote_char)s')',
                aws_commons.create_s3_uri(%(bucket_name)s, %(key)s, %(bucket_region)s)
            );
        """
        params = {"table": table, "bucket_name": bucket_name, "key": key, **kwargs}

        return self.execute_select_query(query, params)