"""
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
    if not db_pool._pool:
        db_pool.initialize()
    
    conn = db_pool._pool.getconn()
    return PooledConnectionWrapper(db_pool._pool, conn)
