from typing import Optional, Any

import sqlparse
from psycopg import sql
from psycopg_pool import AsyncConnectionPool, PoolTimeout

from core.b64_helper import B64Helper
from core.database.database_settings import DatabaseSettings
from core.logger import DBLogger

logger = DBLogger(__name__)


class PostgresPool:
    """
    Create a pool when connecting to PostgreSQL, which will decrease the time spent in
    requesting connection, creating connection, and closing connection.
    """

    def __init__(self, database_settings: DatabaseSettings, pool_size: int):
        self._conninfo = (
            f"host={database_settings.host} "
            f"port={database_settings.port} "
            f"user={database_settings.user} "
            f"password={database_settings.password} "
            f"dbname={database_settings.database}"
        )
        self._pool_size = pool_size

        self.pool: Optional[AsyncConnectionPool] = None

    async def _get_pool(self):
        pool = AsyncConnectionPool(
            conninfo=self._conninfo, open=False, min_size=1, max_size=self._pool_size
        )
        await pool.open()
        try:
            async with pool.connection() as con:
                await pool.check_connection(con)
        except PoolTimeout as e:
            await pool.close()
            logger.fatal(f"Failed to connect to the database", e)
        return pool

    @staticmethod
    async def _exec_sql(cursor: Any, query: str, args=None, multi=True):
        if multi:
            queries = [
                str(stmt).strip() for stmt in sqlparse.parse(query) if str(stmt).strip()
            ]
            for q in queries:
                if q.strip() == "":
                    continue

                await cursor.execute(sql.SQL(q), args)
        else:
            await cursor.execute(sql.SQL(query), args)

    async def execute(self, query: str, args=None, multi=True) -> list[list]:
        """
        Execute a SQL statement, it could be with args and without args. The usage is
        similar to the execute() function in the psycopg module.
        :param query: SQL clause
        :param args: args needed by the SQL clause
        :param multi: if the query is a multi-statement
        :return: return result
        """
        async with await self._get_pool() as pool:
            async with pool.connection() as con:
                async with con.cursor() as cursor:
                    await self._exec_sql(cursor, query, args, multi)

                    if (
                        cursor.description is not None
                    ):  # Check if the query returns rows
                        res = await cursor.fetchall()
                        if res is None:
                            return []

                        result = []
                        for row in res:
                            result.append(list(row))
                        return result
                    else:
                        return []

    async def select(self, query: str, args=None, multi=True) -> list[str]:
        """
        Execute a SQL statement, it could be with args and without args. The usage is
        similar to the execute() function in the psycopg module.
        :param query: SQL clause
        :param args: args needed by the SQL clause
        :param multi: if the query is a multi-statement
        :return: return result
        """
        async with await self._get_pool() as pool:
            async with pool.connection() as con:
                async with con.cursor() as cursor:
                    await self._exec_sql(cursor, query, args, multi)

                    res = await cursor.fetchall()
                    return list(res)

    async def select_map(self, query: str, args=None, multi=True) -> list[dict]:
        """
        Execute a SQL statement, it could be with args and without args. The usage is
        similar to the execute() function in the psycopg module.
        :param query: SQL clause
        :param args: args needed by the SQL clause
        :param multi: if the query is a multi-statement
        :return: return result
        """
        async with await self._get_pool() as pool:
            async with pool.connection() as con:
                async with con.cursor() as cursor:
                    await self._exec_sql(cursor, query, args, multi)

                    res = await cursor.fetchall()
                    res_map: list[dict] = []

                    for i_res in range(len(res)):
                        cols = {}
                        for i_col in range(len(res[i_res])):
                            cols[cursor.description[i_col].name] = res[i_res][i_col]

                        res_map.append(cols)

                    return res_map
