import asyncio
import os
from contextlib import asynccontextmanager

import asyncpg
from dotenv import load_dotenv

# Global connection pool
_pool = None
load_dotenv()


async def get_pool(schema="talmudexplore", min_size=2, max_size=5):
    """Initialize and return the connection pool with the specified schema."""
    global _pool
    if _pool is not None:
        current_loop = asyncio.get_running_loop()
        if getattr(_pool, '_loop', None) != current_loop:
            try:
                await _pool.close()
            except:
                pass
            _pool = None

    if _pool is None:
        _pool = await asyncpg.create_pool(
            database=os.getenv("pg_dbname"),
            user=os.getenv("pg_user"),
            password=os.getenv("pg_password"),
            host=os.getenv("pg_host"),
            port=os.getenv("pg_port"),
            min_size=min_size,
            max_size=max_size,
            setup=lambda conn: conn.execute(f'SET search_path TO {schema}')

        )
    return _pool

@asynccontextmanager
async def get_async_connection():
    """Get a connection from the pool as an async context manager."""
    pool = await get_pool()
    conn = await pool.acquire()
    try:
        yield conn
    finally:
        await pool.release(conn)

async def close_pool():
    """Close the connection pool."""
    global _pool
    if _pool:
        await _pool.close()
        _pool = None