app/database.py

from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from app.config import settings
from app.models.database_models import Base, Document, PathologyReport, User


def _normalize_database_url(url: str) -> str:
    """Normalize common Postgres URL variants for SQLAlchemy."""
    url = (url or "").strip()
    # Heroku/Neon-style URLs sometimes use postgres:// which SQLAlchemy doesn't accept.
    if url.startswith("postgres://"):
        return "postgresql://" + url[len("postgres://") :]
    return url


def _build_database_url() -> str:
    """Select DB target.

    Rules:
    - If DATABASE_URL is set -> use it (Neon / cloud Postgres).
    - Else -> build from POSTGRES_* (local fallback).
    """
    if getattr(settings, "DATABASE_URL", ""):
        return _normalize_database_url(settings.DATABASE_URL)

    return (
        f"postgresql://{settings.POSTGRES_USER}:{settings.POSTGRES_PASSWORD}"
        f"@{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/{settings.POSTGRES_DB}"
    )


# Create database URL
DATABASE_URL = _build_database_url()

# Create engine with persistent connection pool to eliminate SSL renegotiation latency
engine = create_engine(
    DATABASE_URL,
    echo=False,
    pool_size=10,
    max_overflow=20,
    pool_pre_ping=True,
    pool_recycle=300,
    connect_args={"connect_timeout": 10}
)

# Create session
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def get_db():
    """Dependency to get database session"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def _auto_migrate_vector_dimension():
    """
    Auto-migration: upgrade the embedding column from vector(384) → vector(768).

    This runs every startup but is a no-op if the column is already 768-dim or
    if the table does not yet exist. No manual migration script needed.
    """
    TARGET_DIM = 768
    try:
        with engine.connect() as conn:
            # 1. Check table exists
            table_exists = conn.execute(
                text(
                    "SELECT EXISTS ("
                    "  SELECT FROM information_schema.tables "
                    "  WHERE table_name = 'document_embeddings'"
                    ")"
                )
            ).scalar()
            if not table_exists:
                return  # Table will be created fresh with correct dim by create_all()

            # 2. Read current dimension from pg_attribute
            row = conn.execute(
                text(
                    "SELECT atttypmod FROM pg_attribute "
                    "JOIN pg_class ON pg_class.oid = pg_attribute.attrelid "
                    "WHERE pg_class.relname = 'document_embeddings' "
                    "  AND pg_attribute.attname = 'embedding'"
                )
            ).fetchone()

            if row is None:
                return  # No embedding column yet

            current_dim = row[0]  # atttypmod stores the dimension for vector type
            if current_dim == TARGET_DIM:
                return  # Already correct — nothing to do

            # 3. Dimension mismatch — resize column (drop + recreate, pgvector limitation)
            print(
                f"[INFO] Auto-migrating embedding column: {current_dim}-dim -> {TARGET_DIM}-dim "
                f"(BioLORD-2023-M upgrade). Existing vectors cleared."
            )
            conn.execute(
                text("ALTER TABLE document_embeddings DROP COLUMN IF EXISTS embedding")
            )
            conn.execute(
                text(f"ALTER TABLE document_embeddings ADD COLUMN embedding vector({TARGET_DIM})")
            )
            conn.commit()
            print(f"[INFO] Embedding column upgraded to vector({TARGET_DIM}) [OK]")

    except Exception as e:
        # Non-fatal — log and continue. The app still works; embeddings just won't save.
        print(f"[WARN] Vector dimension auto-migration failed (non-fatal): {e}")


def init_db():
    """Initialize database tables"""
    try:
        # Required for pgvector columns (DocumentEmbedding.embedding).
        # On managed Postgres (e.g., Neon), this typically works, but some roles may not
        # have privileges. If extension creation fails, we still create the core tables
        # needed for Phase 1–6 persistence (documents + pathology_reports + users).
        vector_enabled = False
        try:
            with engine.begin() as conn:
                conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
            vector_enabled = True
        except Exception as ext_err:
            print(
                "[WARN] pgvector extension could not be enabled. "
                "Skipping vector tables (document_embeddings). "
                f"Reason: {ext_err}"
            )

        if vector_enabled:
            Base.metadata.create_all(bind=engine)
            # Auto-migrate vector dimension for BioLORD-2023-M (384 → 768)
            _auto_migrate_vector_dimension()
        else:
            Base.metadata.create_all(
                bind=engine,
                tables=[Document.__table__, PathologyReport.__table__, User.__table__],
            )
        return True
    except Exception as e:
        print(f"Database initialization error: {e}")
        return False


# Create tables on startup/import
init_db()