# db/pool.py
"""
PostgreSQL connection pool management.

This module provides:
1. Global pool for NoETL's internal database (server use)
2. Per-credential pool registry for plugin executions (worker use)
"""
from psycopg import AsyncConnection
from psycopg_pool import AsyncConnectionPool
from psycopg.rows import dict_row, DictRow
import asyncio
import hashlib
from typing import Optional, Dict
from contextlib import asynccontextmanager
from noetl.core.logger import setup_logger

logger = setup_logger(__name__, include_location=True)

# Global pool for NoETL's internal database
_pool: Optional[AsyncConnectionPool[AsyncConnection[DictRow]]] = None
_lock = asyncio.Lock()


async def init_pool(conninfo: str):
    """
    Initialize the global AsyncConnectionPool with dict_row as default.
    Safe to call multiple times.
    """
    global _pool
    async with _lock:
        if _pool is None:
            _pool = AsyncConnectionPool(
                conninfo,
                min_size=1,
                max_size=5,
                timeout=10,
                max_lifetime=1800.0,
                max_idle=300.0,
                kwargs={"row_factory": dict_row},
                name="noetl_server",
                open=False
            )
            await _pool.open(wait=True)


def get_pool() -> AsyncConnectionPool[AsyncConnection[DictRow]]:
    """Return AsyncConnectionPool with dict_row as default pool instance."""
    global _pool
    if _pool is None:
        raise RuntimeError("Database pool is not initialized. Call init_pool() first.")
    return _pool

@asynccontextmanager
async def get_pool_connection():
    """Get a connection from the global pool as an async context manager."""
    async with get_pool().connection() as conn:
        yield conn


async def get_snowflake_id() -> int:
    """
    Generate a new snowflake ID using the database function noetl.snowflake_id().
    
    This ensures consistent ID generation across the system and allows
    the execution_id to be cached/shared for evaluation purposes.
    
    Returns:
        int: A unique snowflake ID generated by the database
    """
    async with get_pool_connection() as conn:
        async with conn.cursor() as cur:
            await cur.execute("SELECT noetl.snowflake_id()")
            result = await cur.fetchone()
            if result:
                return int(result['snowflake_id'])
            raise RuntimeError("Failed to generate snowflake ID from database")


async def close_pool():
    """Close and reset the global connection pool."""
    global _pool
    async with _lock:
        if _pool is not None:
            await _pool.close()
            _pool = None

