"""Database connection utilities for Quber RAG system."""

from contextlib import contextmanager
from typing import Generator

from loguru import logger
from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker

from quber.db.models import Base
from quber.settings import get_settings


def get_database_url() -> str:
    """
    Construct database URL from `get_settings().db`.

    Each POSTGRES_* variable that is unset falls back to its default in
    `DBSettings`, so a missing variable never raises here.

    Returns:
        Database connection URL string.
    """
    db_settings = get_settings().db
    user = db_settings.user
    password = db_settings.password
    host = db_settings.host
    port = db_settings.port
    db = db_settings.db

    url = f"postgresql+psycopg://{user}:{password}@{host}:{port}/{db}"
    logger.debug(f"Database URL: postgresql+psycopg://{user}:***@{host}:{port}/{db}")
    return url


def get_engine(echo: bool = False) -> Engine:
    """
    Create and return a SQLAlchemy engine.

    Args:
        echo: If True, log all SQL statements (useful for debugging).

    Returns:
        SQLAlchemy Engine instance.
    """
    url = get_database_url()
    engine = create_engine(url, echo=echo, pool_pre_ping=True)
    logger.info(f"Database engine created: {engine.url.database}@{engine.url.host}")
    return engine


def init_db(engine: Engine | None = None, drop_all: bool = False) -> None:
    """
    Initialize database schema.

    Args:
        engine: SQLAlchemy engine. If None, creates a new one.
        drop_all: If True, drop all tables before creating (WARNING: data loss).
    """
    if engine is None:
        engine = get_engine()

    if drop_all:
        logger.warning("Dropping all tables...")
        Base.metadata.drop_all(engine)

    logger.info("Creating database tables...")
    Base.metadata.create_all(engine)
    logger.success("Database initialized successfully")


def get_session_factory(engine: Engine | None = None) -> sessionmaker[Session]:
    """
    Create a session factory.

    Args:
        engine: SQLAlchemy engine. If None, creates a new one.

    Returns:
        SQLAlchemy sessionmaker.
    """
    if engine is None:
        engine = get_engine()

    return sessionmaker(bind=engine, expire_on_commit=False)


@contextmanager
def get_session(engine: Engine | None = None) -> Generator[Session, None, None]:
    """
    Context manager for database sessions.

    Args:
        engine: SQLAlchemy engine. If None, creates a new one.

    Yields:
        Database session.

    Example:
        ```python
        with get_session() as session:
            documents = session.query(Document).all()
        ```
    """
    if engine is None:
        engine = get_engine()

    SessionLocal = sessionmaker(bind=engine, expire_on_commit=False)
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except Exception as e:
        session.rollback()
        logger.error(f"Session error: {e}")
        raise
    finally:
        session.close()