# db/connection.py
import psycopg2
import os
from contextlib import contextmanager
from typing import Generator

from utils.loadDotEnv import initializeENV

initializeENV()

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()
        print("Подключение к БД установлено")
        yield conn
    except psycopg2.OperationalError as e:
        print(f"Ошибка подключения к БД: {e}")
        raise
    except Exception as e:
        print(f"Неожиданная ошибка: {e}")
        raise
    finally:
        if conn:
            conn.close()
            print("Соединение с БД закрыто")

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

print(test_connection())