import psycopg2
import os
import logging
from contextlib import contextmanager
from typing import Generator

from utils.loadDotEnv import initializeENV

initializeENV()

logger = logging.getLogger(__name__)

def PSQLConnect():
    conn = psycopg2.connect(os.getenv('POSTDRESS_CONNECTION'))
    return conn

def PSQLCursor(conn):
    cur =  conn.cursor()
    return cur


@contextmanager
def get_connection() -> Generator[psycopg2.extensions.connection, None, None]:
    conn = None
    try:
        conn = PSQLConnect()
        logger.info("Подключение к БД установлено")
        yield conn
    except psycopg2.OperationalError as e:
        logger.error(f"Ошибка подключения к БД: {e}")
        raise
    except psycopg2.Error as e:
        logger.error(f"Ошибка PostgreSQL: {e}")
        raise
    except Exception as e:
        logger.error(f"Неожиданная ошибка при работе с БД: {e}")
        raise
    finally:
        if conn:
            try:
                conn.close()
                logger.info("Соединение с БД закрыто")
            except Exception as e:
                logger.warning(f"Ошибка при закрытии соединения: {e}")

def test_connection() -> bool:
    try:
        with get_connection() as conn:
            cur = PSQLCursor(conn)
            try:
                cur.execute("SELECT version();")
                version = cur.fetchone()
                logger.info(f"Версия PostgreSQL: {version[0]}")
                return True
            finally:
                cur.close()
                logger.debug("Курсор закрыт")
    except Exception as e:
        logger.error(f"Тест подключения к БД провален: {e}")
        return False
print(test_connection())