"""PostgreSQL database connection module"""
import psycopg2
import os
from contextlib import contextmanager

DB_CONFIG = {
    "host": os.environ.get("DB_HOST", "pgm-bp1t1008h019ez6c.pg.rds.aliyuncs.com"),
    "port": int(os.environ.get("DB_PORT", 5432)),
    "user": os.environ.get("DB_USER", "coolbot"),
    "password": os.environ.get("DB_PASSWORD", "Coolbot123"),
    "database": os.environ.get("DB_NAME", "coolbot_data"),
}

@contextmanager
def get_db():
    conn = None
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        yield conn
        conn.commit()
    except Exception as e:
        if conn:
            conn.rollback()
        raise e
    finally:
        if conn:
            conn.close()

class DictCursor:
    def __init__(self, cursor):
        self._cursor = cursor

    def __iter__(self):
        columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
        for row in self._cursor:
            yield dict(zip(columns, row))

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    def execute(self, *args, **kwargs):
        return self._cursor.execute(*args, **kwargs)

    def fetchone(self):
        row = self._cursor.fetchone()
        if row is None:
            return None
        columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
        return dict(zip(columns, row))

    def fetchall(self):
        columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
        return [dict(zip(columns, row)) for row in self._cursor.fetchall()]

    def fetchmany(self, size=None):
        if size is None:
            size = self._cursor.arraysize
        columns = [desc[0] for desc in self._cursor.description] if self._cursor.description else []
        rows = self._cursor.fetchmany(size)
        return [dict(zip(columns, row)) for row in rows]

    def close(self):
        return self._cursor.close()

class Database:
    @contextmanager
    def get_cursor(self, dictionary=True):
        conn = None
        try:
            conn = psycopg2.connect(**DB_CONFIG)
            raw_cursor = conn.cursor()
            try:
                if dictionary:
                    cursor = DictCursor(raw_cursor)
                else:
                    cursor = raw_cursor
                yield cursor
                conn.commit()
            except Exception as e:
                conn.rollback()
                raise e
            finally:
                raw_cursor.close()
        finally:
            if conn:
                conn.close()

db = Database()

def test_connection():
    try:
        with get_db() as conn:
            cur = conn.cursor()
            cur.execute("SELECT version()")
            print(f"DB OK: {str(cur.fetchone()[0])[:50]}")
            return True
    except Exception as e:
        print(f"DB ERROR: {e}")
        return False

if __name__ == "__main__":
    test_connection()