from collections import OrderedDict
from collections.abc import Iterator
from contextlib import contextmanager
from contextvars import ContextVar
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_pgconn_string

_check_connection = ConnectionPool.check_connection
_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
_pool_conninfo: dict[str, str] = {}
_pool_borrows: dict[str, int] = {}
_admin_pools: OrderedDict[str, ConnectionPool] = OrderedDict()
_admin_pool_borrows: dict[str, int] = {}
_registry_lock = RLock()
_active_project_connection: ContextVar[tuple[str, Connection] | None] = ContextVar(
    "wndb_active_project_connection",
    default=None,
)
_active_model_mutation_locks: ContextVar[frozenset[str]] = ContextVar(
    "wndb_active_model_mutation_locks",
    default=frozenset(),
)


def _close_pool(pool: ConnectionPool) -> None:
    if not pool.closed:
        pool.close()


def _evict_idle_project_pools(*, protected: str | None = None) -> None:
    limit = max(1, settings.PROJECT_PG_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)
        _close_pool(pool)


def _evict_idle_admin_pools(*, protected: str | None = None) -> None:
    limit = max(1, settings.PROJECT_PG_CACHE_SIZE)
    while len(_admin_pools) > limit:
        candidate = next(
            (
                key
                for key in _admin_pools
                if key != protected and _admin_pool_borrows.get(key, 0) == 0
            ),
            None,
        )
        if candidate is None:
            return
        pool = _admin_pools.pop(candidate)
        _admin_pool_borrows.pop(candidate, None)
        _close_pool(pool)


def get_project_pool(name: str) -> ConnectionPool:
    """Return the routed synchronous pool used by native WNDB operations."""
    conninfo = get_project_pgconn_string(db_name=name)
    with _registry_lock:
        pool = _pools.get(name)
        if pool is not None and _pool_conninfo.get(name) == conninfo and not pool.closed:
            _pools.move_to_end(name)
            return pool
        if pool is not None:
            if _pool_borrows.get(name, 0):
                raise RuntimeError(f"Cannot replace active project pool {name!r}")
            _close_pool(pool)
        pool = ConnectionPool(
            conninfo=conninfo,
            min_size=settings.PROJECT_PG_POOL_MIN_SIZE,
            max_size=settings.PROJECT_PG_POOL_SIZE + settings.PROJECT_PG_MAX_OVERFLOW,
            kwargs={"autocommit": True, "row_factory": dict_row},
            check=_check_connection,
            open=True,
        )
        _pools[name] = pool
        _pool_conninfo[name] = conninfo
        _pool_borrows.setdefault(name, 0)
        _evict_idle_project_pools(protected=name)
        return pool


def is_project_pool_open(name: str) -> bool:
    with _registry_lock:
        pool = _pools.get(name)
        if pool is None or pool.closed:
            return False
        if _pool_conninfo.get(name) != get_project_pgconn_string(db_name=name):
            if _pool_borrows.get(name, 0):
                return False
            _close_pool(pool)
            _pools.pop(name, None)
            _pool_conninfo.pop(name, None)
            _pool_borrows.pop(name, None)
            return False
        return True


def close_project_pool(name: str) -> None:
    with _registry_lock:
        if _pool_borrows.get(name, 0):
            raise RuntimeError(f"Cannot close active project pool {name!r}")
        pool = _pools.pop(name, None)
        _pool_conninfo.pop(name, None)
        _pool_borrows.pop(name, None)
        if pool is not None:
            _close_pool(pool)


def close_all_project_pools() -> None:
    """Close every WNDB project pool and the administration pool."""
    with _registry_lock:
        pools = [*_pools.values(), *_admin_pools.values()]
        _pools.clear()
        _pool_conninfo.clear()
        _pool_borrows.clear()
        _admin_pools.clear()
        _admin_pool_borrows.clear()
        for pool in pools:
            _close_pool(pool)


def get_admin_pool() -> ConnectionPool:
    """Return the administration pool for the current routed PostgreSQL host."""
    conninfo = get_project_pgconn_string(db_name="postgres")
    with _registry_lock:
        pool = _admin_pools.get(conninfo)
        if pool is not None and not pool.closed:
            _admin_pools.move_to_end(conninfo)
            return pool
        pool = ConnectionPool(
            conninfo=conninfo,
            min_size=settings.PROJECT_PG_POOL_MIN_SIZE,
            max_size=settings.PROJECT_PG_POOL_SIZE,
            kwargs={"autocommit": True, "row_factory": dict_row},
            check=_check_connection,
            open=True,
        )
        _admin_pools[conninfo] = pool
        _admin_pool_borrows.setdefault(conninfo, 0)
        _evict_idle_admin_pools(protected=conninfo)
        return pool


@contextmanager
def project_connection(name: str) -> Iterator[Connection]:
    """Borrow one routed WNDB connection and return it to its pool on exit."""
    active = _active_project_connection.get()
    if active is not None:
        active_name, conn = active
        if active_name != name:
            raise RuntimeError(
                f"Cannot access project {name!r} inside transaction for {active_name!r}"
            )
        yield conn
        return
    with _registry_lock:
        pool = get_project_pool(name)
        _pool_borrows[name] = _pool_borrows.get(name, 0) + 1
    try:
        with pool.connection() as conn:
            yield conn
    finally:
        with _registry_lock:
            _pool_borrows[name] -= 1
            _evict_idle_project_pools()


@contextmanager
def project_transaction(name: str) -> Iterator[Connection]:
    """Run all nested WNDB operations on one pooled connection and transaction."""
    active = _active_project_connection.get()
    if active is not None:
        active_name, conn = active
        if active_name != name:
            raise RuntimeError(
                f"Cannot nest project {name!r} inside transaction for {active_name!r}"
            )
        with conn.transaction():
            yield conn
        return

    with _registry_lock:
        pool = get_project_pool(name)
        _pool_borrows[name] = _pool_borrows.get(name, 0) + 1
    try:
        with pool.connection() as conn:
            token = _active_project_connection.set((name, conn))
            lock_token = _active_model_mutation_locks.set(frozenset())
            try:
                with conn.transaction():
                    yield conn
            finally:
                _active_model_mutation_locks.reset(lock_token)
                _active_project_connection.reset(token)
    finally:
        with _registry_lock:
            _pool_borrows[name] -= 1
            _evict_idle_project_pools()


def is_project_transaction_active(name: str) -> bool:
    active = _active_project_connection.get()
    return active is not None and active[0] == name


def is_model_mutation_lock_active(name: str) -> bool:
    """Return whether the current project transaction already owns its model lock."""
    return name in _active_model_mutation_locks.get()


def mark_model_mutation_lock_active(name: str) -> None:
    """Record a transaction-scoped advisory lock to avoid duplicate round trips."""
    locks = _active_model_mutation_locks.get()
    _active_model_mutation_locks.set(locks | {name})


@contextmanager
def admin_connection() -> Iterator[Connection]:
    """Borrow a PostgreSQL administration connection from its pool."""
    with _registry_lock:
        pool = get_admin_pool()
        conninfo = pool.conninfo
        _admin_pool_borrows[conninfo] = _admin_pool_borrows.get(conninfo, 0) + 1
    try:
        with pool.connection() as conn:
            yield conn
    finally:
        with _registry_lock:
            _admin_pool_borrows[conninfo] -= 1
            _evict_idle_admin_pools()