import psycopg2
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv
import os

load_dotenv()

def get_connection():
    # return psycopg2.connect(
    #     host="localhost",
    #     port=5432,
    #     dbname="wfl_db",
    #     user="postgres",
    #     password="159753"
    # )
    return psycopg2.connect(
        host=os.getenv("DB_HOST"),
        port=os.getenv("DB_PORT"),
        dbname=os.getenv("DB_NAME"),
        user=os.getenv("DB_USER"),
        password=os.getenv("DB_PASSWORD")
    )


def test_connection():
    conn = get_connection()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("SELECT current_database() AS db_name;")
            row = cur.fetchone()
            print(f"Connected to database: {row['db_name']}")
    finally:
        conn.close()


if __name__ == "__main__":
    test_connection()