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

import psycopg as pg

from app.core.config import get_pgconn_string

g_conn_dict: dict[str, pg.Connection] = {}
_registry_lock = RLock()
_project_locks: dict[str, RLock] = {}


def _is_closed(connection: pg.Connection) -> bool:
    return bool(getattr(connection, "closed", False))


def _close_connection(connection: pg.Connection) -> None:
    if not _is_closed(connection):
        connection.close()


def _is_healthy(connection: pg.Connection) -> bool:
    if _is_closed(connection):
        return False
    try:
        with connection.cursor() as cur:
            cur.execute("SELECT 1")
    except pg.Error:
        return False
    return True


def _get_project_lock(name: str) -> RLock:
    with _registry_lock:
        lock = _project_locks.get(name)
        if lock is None:
            lock = RLock()
            _project_locks[name] = lock
        return lock


def open_connection(name: str) -> pg.Connection:
    with _get_project_lock(name):
        connection = g_conn_dict.get(name)
        if connection is None or not _is_healthy(connection):
            if connection is not None:
                _close_connection(connection)
            connection = pg.connect(
                conninfo=get_pgconn_string(db_name=name), autocommit=True
            )
            g_conn_dict[name] = connection
        return connection


def is_connection_open(name: str) -> bool:
    with _get_project_lock(name):
        connection = g_conn_dict.get(name)
        if connection is None:
            return False
        if not _is_healthy(connection):
            del g_conn_dict[name]
            _close_connection(connection)
            return False
        return True


def close_connection(name: str) -> None:
    with _get_project_lock(name):
        connection = g_conn_dict.pop(name, None)
        if connection is not None:
            _close_connection(connection)


@contextmanager
def project_connection(name: str) -> Iterator[pg.Connection]:
    with _get_project_lock(name):
        yield open_connection(name)