from __future__ import annotations

from contextlib import contextmanager
from typing import Iterator

from pgvector.psycopg import register_vector
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool

from app.config import get_settings

_pool: ConnectionPool | None = None


def init_pool() -> ConnectionPool:
    global _pool
    settings = get_settings()
    _pool = ConnectionPool(
        conninfo=settings.database_url,
        min_size=0,
        max_size=8,
        timeout=15,
        kwargs={"row_factory": dict_row, "autocommit": False},
        configure=register_vector,
        open=True,
    )
    return _pool


def close_pool() -> None:
    global _pool
    if _pool is not None:
        _pool.close()
        _pool = None


def get_pool() -> ConnectionPool:
    if _pool is None:
        raise RuntimeError("Database pool is not initialized.")
    return _pool


@contextmanager
def get_connection() -> Iterator:
    pool = get_pool()
    with pool.connection() as conn:
        register_vector(conn)
        yield conn