"""PostgreSQL connection pool management for monitoring.

Uses asyncpg for async database access.
"""

from __future__ import annotations

import json
from typing import Any

from fluxibly.exceptions import MonitoringDBError

try:
    import asyncpg
except ImportError as e:
    raise ImportError(
        "asyncpg is required for monitoring. "
        "Install it with: pip install fluxibly[monitoring]"
    ) from e


# Module-level pool cache keyed by DSN
_pools: dict[str, asyncpg.Pool] = {}


async def _init_connection(conn: asyncpg.Connection) -> None:
    """Register JSONB/JSON codecs so values are auto-deserialized to Python objects."""
    await conn.set_type_codec(
        "jsonb",
        encoder=lambda v: v,  # _json_dumps already provides JSON strings
        decoder=json.loads,
        schema="pg_catalog",
        format="text",
    )
    await conn.set_type_codec(
        "json",
        encoder=lambda v: v,
        decoder=json.loads,
        schema="pg_catalog",
        format="text",
    )


async def create_pool(
    host: str = "localhost",
    port: int = 5432,
    database: str = "fluxibly_monitoring",
    user: str = "postgres",
    password: str = "",
    min_size: int = 1,
    max_size: int = 5,
) -> asyncpg.Pool:
    """Create an asyncpg connection pool.

    Raises:
        MonitoringDBError: If connection fails.
    """
    dsn = f"postgresql://{user}@{host}:{port}/{database}"

    if dsn in _pools:
        pool = _pools[dsn]
        # Check if pool is still alive
        try:
            async with pool.acquire() as conn:
                await conn.fetchval("SELECT 1")
            return pool
        except Exception:
            # Pool is dead, close properly and recreate
            try:
                await pool.close()
            except Exception:
                pass
            _pools.pop(dsn, None)

    try:
        pool = await asyncpg.create_pool(
            host=host,
            port=port,
            database=database,
            user=user,
            password=password,
            min_size=min_size,
            max_size=max_size,
            init=_init_connection,
        )
        _pools[dsn] = pool
        return pool
    except Exception as e:
        raise MonitoringDBError(
            f"Failed to create monitoring DB pool: {e}"
        ) from e


async def get_pool(connect_kwargs: dict[str, Any]) -> asyncpg.Pool:
    """Get or create a pool from connection kwargs."""
    return await create_pool(**connect_kwargs)


async def close_all_pools() -> None:
    """Close all cached connection pools."""
    for dsn, pool in list(_pools.items()):
        try:
            await pool.close()
        except Exception:
            pass
    _pools.clear()


async def run_migration(pool: asyncpg.Pool, sql_path: str) -> None:
    """Run a SQL migration file against the pool.

    Args:
        pool: asyncpg connection pool.
        sql_path: Path to the .sql file.

    Raises:
        MonitoringDBError: If migration fails.
    """
    from pathlib import Path

    path = Path(sql_path)
    if not path.exists():
        raise MonitoringDBError(f"Migration file not found: {sql_path}")

    sql = path.read_text()
    try:
        async with pool.acquire() as conn:
            await conn.execute(sql)
    except Exception as e:
        raise MonitoringDBError(f"Migration failed: {e}") from e
