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

from utils.loadDotEnv import initializeENV

initializeENV()


def get_db_config() -> dict:
    return {
        'host': os.getenv('DB_HOST', 'localhost'),
        'port': int(os.getenv('DB_PORT', 5432)),
        'database': os.getenv('DB_NAME', 'warehouse_db'),
        'user': os.getenv('DB_USER', 'postgres'),
        'password': os.getenv('DB_PASSWORD', '')
    }


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


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


print(test_connection())