"""
Wrapper that intercepts .close() to return the connection to the pool instead.
"""
from typing import Any
import psycopg


class PooledConnectionWrapper:
    def __init__(self, pool, conn: psycopg.Connection):
        self._pool = pool
        self._conn = conn
        self._closed = False

    def __getattr__(self, name: str) -> Any:
        return getattr(self._conn, name)

    def close(self):
        if not self._closed:
            self._pool.putconn(self._conn)
            self._closed = True

    def __enter__(self):
        return self._conn.__enter__()
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            self._conn.__exit__(exc_type, exc_val, exc_tb)
        finally:
            self.close()

def get_pooled_connection_compat():
    """Drop-in replacement for psycopg.connect() that returns to the pool on .close()"""
    from common.db_pool import db_pool
    import psycopg
    
    if not db_pool._pool:
        db_pool.initialize()
    
    conn = None
    for _ in range(3):
        try:
            conn = db_pool._pool.getconn()
            # Perform a ping to ensure the TCP connection is still alive (handles AWS/Azure idle drop)
            conn.execute("SELECT 1")
            return PooledConnectionWrapper(db_pool._pool, conn)
        except (psycopg.OperationalError, Exception) as e:
            # Drop bad connection instead of returning it to pool
            if conn:
                try:
                    conn.close() # Close it hard
                except Exception:
                    pass
    
    # Final attempt if all retries hit a broken connection
    conn = db_pool._pool.getconn()
    return PooledConnectionWrapper(db_pool._pool, conn)
