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

import psycopg2
from psycopg2 import sql
from psycopg2.extras import RealDictCursor

from config import POSTGRES_DB, POSTGRES_HOST, POSTGRES_PASSWORD, POSTGRES_PORT, POSTGRES_USER
from db.db_connector import DBConnector


__all__ = ["PostgresConnector", "PostgresConnectionError"]


class PostgresConnectionError(Exception):
    pass


class PostgresConnector(DBConnector):
    user = POSTGRES_USER
    password = POSTGRES_PASSWORD
    host = POSTGRES_HOST
    port = POSTGRES_PORT
    database = POSTGRES_DB

    def close(self):
        self._conn.close()

    def _get_connection(self):
        try:
            return psycopg2.connect(
                user=self.user,
                password=self.password,
                host=self.host,
                port=self.port,
                database=self.database,
                cursor_factory=RealDictCursor,
            )
        except psycopg2.DatabaseError as e:
            raise PostgresConnectionError("Can't connect to postgres DB") from e

    @contextmanager
    def transaction(self) -> RealDictCursor:
        with self._conn:
            with self._conn.cursor() as cursor:
                yield cursor

    def execute_query(self, query: Union[str, sql.Composable], params: Optional[Dict[str, Any]] = None):
        """
        Execute postgresql query
        :param query: Raw SQL string
        :param params: SQL query parameters
        :return: list of dicts
        """
        with self.transaction() as cursor:
            cursor.execute(query, params)

    def execute_select_query(
        self, query: Union[str, sql.Composable], params: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """
        Execute postgresql query and fetch the result
        :param query: Raw SQL string
        :param params: SQL query parameters
        :return: list of dicts
        """
        with self.transaction() as cursor:
            cursor.execute(query, params)
            result = cursor.fetchall()
        return result

    def check_db_connection(self) -> bool:
        results = self.execute_select_query("SELECT 1")
        return results == [{"?column?": 1}]