from collections import OrderedDict
from collections.abc import Iterator
from contextlib import contextmanager
from threading import RLock

from psycopg import Connection
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool

from app.core.config import settings
from app.infra.db.project_routing import get_project_timescale_pgconn_string


_check_connection = ConnectionPool.check_connection
_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
_pool_conninfo: dict[str, str] = {}
_pool_borrows: dict[str, int] = {}
_lock = RLock()


def _evict_idle_pools(*, protected: str | None = None) -> None:
    limit = max(1, settings.PROJECT_TS_CACHE_SIZE)
    while len(_pools) > limit:
        candidate = next(
            (key for key in _pools if key != protected and _pool_borrows.get(key, 0) == 0),
            None,
        )
        if candidate is None:
            return
        pool = _pools.pop(candidate)
        _pool_conninfo.pop(candidate, None)
        _pool_borrows.pop(candidate, None)
        if not pool.closed:
            pool.close()


def get_timescale_pool(db_name: str) -> ConnectionPool:
    conninfo = get_project_timescale_pgconn_string(db_name=db_name)
    with _lock:
        pool = _pools.get(db_name)
        if pool is not None and _pool_conninfo.get(db_name) == conninfo and not pool.closed:
            _pools.move_to_end(db_name)
            return pool
        if pool is not None and not pool.closed:
            if _pool_borrows.get(db_name, 0):
                raise RuntimeError(f"Cannot replace active TimescaleDB pool {db_name!r}")
            pool.close()
        pool = ConnectionPool(
            conninfo=conninfo,
            min_size=settings.PROJECT_TS_POOL_MIN_SIZE,
            max_size=settings.PROJECT_TS_POOL_MAX_SIZE,
            kwargs={"row_factory": dict_row},
            check=_check_connection,
            open=True,
        )
        _pools[db_name] = pool
        _pool_conninfo[db_name] = conninfo
        _pool_borrows.setdefault(db_name, 0)
        _evict_idle_pools(protected=db_name)
        return pool


@contextmanager
def timescale_connection(db_name: str) -> Iterator[Connection]:
    with _lock:
        pool = get_timescale_pool(db_name)
        _pool_borrows[db_name] = _pool_borrows.get(db_name, 0) + 1
    try:
        with pool.connection() as conn:
            yield conn
    finally:
        with _lock:
            _pool_borrows[db_name] -= 1
            _evict_idle_pools()


def close_timescale_pool(db_name: str) -> None:
    with _lock:
        if _pool_borrows.get(db_name, 0):
            raise RuntimeError(f"Cannot close active TimescaleDB pool {db_name!r}")
        pool = _pools.pop(db_name, None)
        _pool_conninfo.pop(db_name, None)
        _pool_borrows.pop(db_name, None)
        if pool is not None and not pool.closed:
            pool.close()


def close_all_timescale_pools() -> None:
    """Close every synchronous TimescaleDB pool."""
    with _lock:
        pools = list(_pools.values())
        _pools.clear()
        _pool_conninfo.clear()
        _pool_borrows.clear()
        for pool in pools:
            if not pool.closed:
                pool.close()