backend/database.py

from sqlalchemy import create_engine, event, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import logging

from backend.config import settings

logger = logging.getLogger(__name__)

IS_SQLITE = "sqlite" in settings.DATABASE_URL

engine_kwargs = {"pool_pre_ping": True, "echo": settings.DEBUG}
if IS_SQLITE:
    engine_kwargs["connect_args"] = {"check_same_thread": False}
else:
    engine_kwargs["pool_size"] = 5
    engine_kwargs["max_overflow"] = 10
    engine_kwargs["pool_timeout"] = 30

engine = create_engine(settings.DATABASE_URL, **engine_kwargs)

if IS_SQLITE:
    @event.listens_for(engine, "connect")
    def set_sqlite_pragma(dbapi_connection, connection_record):
        cursor = dbapi_connection.cursor()
        cursor.execute("PRAGMA journal_mode=WAL")
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.close()

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def _table_has_column(conn, table_name: str, column_name: str) -> bool:
    """Check if a column exists in a PostgreSQL table."""
    result = conn.execute(text("""
        SELECT COUNT(*) FROM information_schema.columns
        WHERE table_name = :t AND column_name = :c
    """), {"t": table_name, "c": column_name})
    return result.scalar() > 0


def create_tables():
    """
    Create all tables safely.
    If an existing table is missing columns (broken schema from old run),
    drop it and recreate cleanly.
    """
    from backend import models  # noqa: F401

    if IS_SQLITE:
        Base.metadata.create_all(bind=engine, checkfirst=True)
        logger.info("✅ SQLite tables ready")
        return

    # PostgreSQL: check if tables exist with correct schema
    with engine.connect() as conn:
        # Check if users table exists but is broken (missing 'username' column)
        result = conn.execute(text(
            "SELECT COUNT(*) FROM information_schema.tables WHERE table_name='users'"
        ))
        users_exists = result.scalar() > 0

        if users_exists and not _table_has_column(conn, "users", "username"):
            logger.warning("⚠️  Broken schema detected — dropping all tables and recreating...")
            conn.execute(text("DROP TABLE IF EXISTS messages CASCADE"))
            conn.execute(text("DROP TABLE IF EXISTS conversations CASCADE"))
            conn.execute(text("DROP TABLE IF EXISTS user_api_keys CASCADE"))
            conn.execute(text("DROP TABLE IF EXISTS users CASCADE"))
            conn.commit()
            logger.info("🗑  Old tables dropped")

    # Now create fresh
    Base.metadata.create_all(bind=engine, checkfirst=True)
    logger.info("✅ PostgreSQL tables ready")