"""Postgres/pgvector connection helper for the ADE RAG playground.

Reuses `quber.settings` for credentials (it loads `.env` from the current
working directory) and registers the pgvector type adapters so `VECTOR`
columns round-trip as Python lists / numpy arrays.
"""

from __future__ import annotations

from contextlib import asynccontextmanager, contextmanager
from pathlib import Path
from typing import AsyncGenerator, Generator, LiteralString, cast

import psycopg
from pgvector.psycopg import register_vector, register_vector_async

from quber.settings import get_settings

SCHEMA_FILE = Path(__file__).with_name("schema.sql")


def dsn() -> str:
    db = get_settings().db
    return f"host={db.host} port={db.port} dbname={db.db} user={db.user} password={db.password}"


@contextmanager
def connect() -> Generator[psycopg.Connection]:
    """Yield a connection with pgvector adapters registered and autocommit on.

    For code off the event loop: ingest, the CLI, scripts. A request handler
    or anything else running under asyncio uses `connect_async`, because a
    blocking query there holds the loop and stalls every other request's
    in-flight work.
    """
    with psycopg.connect(dsn(), autocommit=True) as conn:
        register_vector(conn)
        yield conn


@asynccontextmanager
async def connect_async() -> AsyncGenerator[psycopg.AsyncConnection]:
    """The same connection for code on the event loop: queries are awaited,
    so the loop keeps serving other requests while Postgres works."""
    async with await psycopg.AsyncConnection.connect(dsn(), autocommit=True) as conn:
        await register_vector_async(conn)
        yield conn


def apply_schema() -> None:
    """Create the `ade_playground` schema from schema.sql (drops existing tables)."""
    # psycopg types `execute` to LiteralString to keep injection out of query
    # strings; a schema file we author and read off disk is the same trust level.
    sql = cast(LiteralString, SCHEMA_FILE.read_text())
    with connect() as conn:
        conn.execute(sql)