# 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 hashlib
import os
import threading
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
_pool_lock = threading.Lock()


def _env_int(name: str, default: int) -> int:
    raw = os.getenv(name)
    if raw is None or raw.strip() == "":
        return default
    try:
        return int(raw)
    except (TypeError, ValueError):
        return default


def _env_float(name: str, default: float) -> float:
    raw = os.getenv(name)
    if raw is None or raw.strip() == "":
        return default
    try:
        return float(raw)
    except (TypeError, ValueError):
        return default


_DEFAULT_POOL_MIN_SIZE = max(1, _env_int("NOETL_POSTGRES_POOL_MIN_SIZE", 2))
_DEFAULT_POOL_MAX_SIZE = max(_DEFAULT_POOL_MIN_SIZE, _env_int("NOETL_POSTGRES_POOL_MAX_SIZE", 16))
_DEFAULT_POOL_TIMEOUT = max(1.0, _env_float("NOETL_POSTGRES_POOL_TIMEOUT_SECONDS", 30.0))
_DEFAULT_POOL_MAX_WAITING = max(1, _env_int("NOETL_POSTGRES_POOL_MAX_WAITING", 256))
_DEFAULT_POOL_MAX_LIFETIME = max(60.0, _env_float("NOETL_POSTGRES_POOL_MAX_LIFETIME_SECONDS", 1800.0))
_DEFAULT_POOL_MAX_IDLE = max(30.0, _env_float("NOETL_POSTGRES_POOL_MAX_IDLE_SECONDS", 300.0))
_DEFAULT_STATEMENT_TIMEOUT_MS = max(0, _env_int("NOETL_POSTGRES_STATEMENT_TIMEOUT_MS", 60000))
_DEFAULT_IDLE_IN_TX_TIMEOUT_MS = max(
    0,
    _env_int("NOETL_POSTGRES_IDLE_IN_TRANSACTION_SESSION_TIMEOUT_MS", 45000),
)


def _build_pool_connect_kwargs() -> dict:
    kwargs = {"row_factory": dict_row}
    options: list[str] = []
    if _DEFAULT_STATEMENT_TIMEOUT_MS > 0:
        options.append(f"-c statement_timeout={_DEFAULT_STATEMENT_TIMEOUT_MS}")
    if _DEFAULT_IDLE_IN_TX_TIMEOUT_MS > 0:
        options.append(
            f"-c idle_in_transaction_session_timeout={_DEFAULT_IDLE_IN_TX_TIMEOUT_MS}"
        )
    if options:
        kwargs["options"] = " ".join(options)
    return kwargs


_POOL_CONNECT_KWARGS = _build_pool_connect_kwargs()


async def init_pool(conninfo: str):
    """
    Initialize the global AsyncConnectionPool with dict_row as default.
    Safe to call multiple times.

    Pool sizing is controlled by NOETL_POSTGRES_POOL_* env vars with safe defaults.
    """
    global _pool
    pool_to_open: Optional[AsyncConnectionPool[AsyncConnection[DictRow]]] = None

    # Use a thread lock instead of asyncio.Lock to avoid cross-event-loop binding.
    with _pool_lock:
        if _pool is None:
            logger.info(
                "DB.POOL: init min=%d max=%d waiting=%d timeout=%.1fs stmt_timeout_ms=%d idle_tx_timeout_ms=%d",
                _DEFAULT_POOL_MIN_SIZE,
                _DEFAULT_POOL_MAX_SIZE,
                _DEFAULT_POOL_MAX_WAITING,
                _DEFAULT_POOL_TIMEOUT,
                _DEFAULT_STATEMENT_TIMEOUT_MS,
                _DEFAULT_IDLE_IN_TX_TIMEOUT_MS,
            )
            _pool = AsyncConnectionPool(
                conninfo,
                min_size=_DEFAULT_POOL_MIN_SIZE,
                max_size=_DEFAULT_POOL_MAX_SIZE,
                max_waiting=_DEFAULT_POOL_MAX_WAITING,
                timeout=_DEFAULT_POOL_TIMEOUT,
                max_lifetime=_DEFAULT_POOL_MAX_LIFETIME,
                max_idle=_DEFAULT_POOL_MAX_IDLE,
                kwargs=_POOL_CONNECT_KWARGS,
                name=os.getenv("NOETL_POSTGRES_POOL_NAME", "noetl_server"),
                open=False,
            )
            pool_to_open = _pool

    if pool_to_open is not None:
        await pool_to_open.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(timeout: float | None = None):
    """Get a connection from the global pool as an async context manager."""
    pool = get_pool()
    if timeout is None:
        async with pool.connection() as conn:
            yield conn
    else:
        async with pool.connection(timeout=timeout) 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")


def get_server_pool_stats() -> dict:
    """
    Return real-time stats for the server's global connection pool.

    Uses psycopg_pool's get_stats() which returns:
      pool_min, pool_max, pool_size, pool_available,
      requests_waiting, requests_made, usage_ms, ...

    Returns an empty dict if the pool has not been initialized yet.
    """
    global _pool
    if _pool is None:
        return {}
    try:
        raw = _pool.get_stats()
        pool_max = int(raw.get("pool_max") or _DEFAULT_POOL_MAX_SIZE)
        pool_available = int(raw.get("pool_available") or 0)
        pool_size = int(raw.get("pool_size") or 0)
        requests_waiting = int(raw.get("requests_waiting") or 0)
        utilization = round(
            (pool_size - pool_available) / pool_max if pool_max > 0 else 0.0, 4
        )
        return {
            "pool_min": int(raw.get("pool_min") or _DEFAULT_POOL_MIN_SIZE),
            "pool_max": pool_max,
            "pool_size": pool_size,
            "pool_available": pool_available,
            "requests_waiting": requests_waiting,
            "utilization": utilization,
            # Derived: how many more connections can be handed out right now
            "slots_available": max(0, pool_available),
        }
    except Exception:
        return {}


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