"""Shared database helpers."""

from __future__ import annotations

import threading
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path

from sqlalchemy import create_engine, text
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker


class Base(DeclarativeBase):
    """Declarative ORM base class."""


def create_engine_with_wal(db_url: str, echo: bool = False):
    """Create SQLAlchemy engine and enable SQLite WAL pragmas when applicable."""
    engine = create_engine(db_url, echo=echo, future=True)
    if db_url.startswith("sqlite"):
        with engine.begin() as conn:
            conn.execute(text("PRAGMA journal_mode=WAL;"))
            conn.execute(text("PRAGMA foreign_keys=ON;"))
            conn.execute(text("PRAGMA busy_timeout=5000;"))
    return engine


def create_session_factory(engine) -> sessionmaker[Session]:
    """Create a SQLAlchemy session factory."""
    return sessionmaker(bind=engine, autoflush=False, autocommit=False, expire_on_commit=False, future=True)


@contextmanager
def session_scope(session_factory: sessionmaker[Session]) -> Iterator[Session]:
    """Context helper for transaction scoping."""
    session = session_factory()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


# ---------------------------------------------------------------------------
# ChromaDB client factory (singleton per persist directory)
# ---------------------------------------------------------------------------

_chroma_clients: dict[str, object] = {}
_chroma_lock = threading.Lock()


def get_chroma_client(persist_dir: str | Path | None = None):
    """Return a persistent ChromaDB client (singleton per directory).

    If *persist_dir* is ``None``, an ephemeral in-memory client is returned
    (useful for tests).
    """
    import chromadb  # lazy — avoid import cost when not needed

    key = str(persist_dir) if persist_dir else ":memory:"
    with _chroma_lock:
        if key not in _chroma_clients:
            if persist_dir:
                path = Path(persist_dir)
                path.mkdir(parents=True, exist_ok=True)
                _chroma_clients[key] = chromadb.PersistentClient(path=str(path))
            else:
                _chroma_clients[key] = chromadb.EphemeralClient()
        return _chroma_clients[key]


def get_or_create_chroma_collection(
    persist_dir: str | Path | None,
    name: str,
    *,
    distance_fn: str = "cosine",
):
    """Get or create a ChromaDB collection by name.

    *distance_fn* can be ``"cosine"``, ``"l2"``, or ``"ip"``.
    """
    client = get_chroma_client(persist_dir)
    return client.get_or_create_collection(name=name, metadata={"hnsw:space": distance_fn})
