backend/app/database.py

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, DeclarativeBase
from sqlalchemy.pool import QueuePool, StaticPool
from app.config import settings

# Handle SQLite vs PostgreSQL connection args
connect_args = {}
engine_kwargs = {}

if settings.DATABASE_URL.startswith("sqlite"):
    connect_args["check_same_thread"] = False
    engine_kwargs["poolclass"] = StaticPool
else:
    # PostgreSQL connection pool settings
    engine_kwargs["poolclass"] = QueuePool
    engine_kwargs["pool_size"] = 5
    engine_kwargs["max_overflow"] = 10
    engine_kwargs["pool_pre_ping"] = True  # reconnect on stale connections

    # Neon requires SSL
    if "sslmode" not in settings.DATABASE_URL:
        connect_args["sslmode"] = "require"

engine = create_engine(
    settings.DATABASE_URL,
    connect_args=connect_args,
    **engine_kwargs,
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


class Base(DeclarativeBase):
    pass


def get_db():
    """FastAPI dependency that provides a database session."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def init_db():
    """Create all tables, run lightweight migrations, and seed plans."""
    from app.models import User, Project, Scene, Asset, ChatMessage, SubscriptionPlan, Subscription  # noqa: F401

    Base.metadata.create_all(bind=engine)
    _migrate(engine)

    # Seed subscription plans on every startup (idempotent)
    from app.models.subscription import seed_plans
    db = SessionLocal()
    try:
        seed_plans(db)
    finally:
        db.close()


def _migrate(eng):
    """Add columns that may be missing from older schemas."""
    from sqlalchemy import text, inspect

    insp = inspect(eng)
    if "projects" not in insp.get_table_names():
        return

    cols = {c["name"] for c in insp.get_columns("projects")}
    is_pg = not settings.DATABASE_URL.startswith("sqlite")

    with eng.begin() as conn:
        # Helper: ALTER TABLE ADD COLUMN (Postgres ignores IF NOT EXISTS style,
        # so we check the column set in Python first)
        migrations = {
            "voice_gender": "VARCHAR(10) DEFAULT 'female'",
            "voice_accent": "VARCHAR(10) DEFAULT 'american'",
            "player_port": "INTEGER",
            "accent_color": "VARCHAR(20) DEFAULT '#7C3AED'",
            "bg_color": "VARCHAR(20) DEFAULT '#FFFFFF'",
            "text_color": "VARCHAR(20) DEFAULT '#000000'",
            "animation_instructions": "TEXT",
            "studio_unlocked": "BOOLEAN DEFAULT 0",
            "r2_video_key": "VARCHAR(512)",
            "r2_video_url": "VARCHAR(2048)",
            "logo_r2_key": "VARCHAR(512)",
            "logo_r2_url": "VARCHAR(2048)",
            "logo_position": "VARCHAR(20) DEFAULT 'bottom_right'",
            "custom_voice_id": "VARCHAR(100)",
            "aspect_ratio": "VARCHAR(20) DEFAULT 'landscape'",
            "template": "VARCHAR(50) DEFAULT 'default'",
            "ai_assisted_editing_count": "INTEGER DEFAULT 0",
        }
        for col_name, col_def in migrations.items():
            if col_name not in cols:
                conn.execute(text(
                    f"ALTER TABLE projects ADD COLUMN {col_name} {col_def}"
                ))

    # Migrate users table
    if "users" in insp.get_table_names():
        user_cols = {c["name"] for c in insp.get_columns("users")}
        with eng.begin() as conn:
            user_migrations = {
                "video_limit_bonus": "INTEGER DEFAULT 0",
            }
            for col_name, col_def in user_migrations.items():
                if col_name not in user_cols:
                    conn.execute(text(
                        f"ALTER TABLE users ADD COLUMN {col_name} {col_def}"
                    ))

    # Migrate assets table
    if "assets" in insp.get_table_names():
        asset_cols = {c["name"] for c in insp.get_columns("assets")}
        with eng.begin() as conn:
            asset_migrations = {
                "r2_key": "VARCHAR(512)",
                "r2_url": "VARCHAR(2048)",
                "excluded": "BOOLEAN DEFAULT 0",
            }
            for col_name, col_def in asset_migrations.items():
                if col_name not in asset_cols:
                    conn.execute(text(
                        f"ALTER TABLE assets ADD COLUMN {col_name} {col_def}"
                    ))

    # Migrate PostgreSQL enum types — add missing values.
    # ALTER TYPE ... ADD VALUE cannot run inside a transaction, so we use
    # a raw DBAPI connection with autocommit.
    if is_pg:
        try:
            raw_conn = eng.raw_connection()
            raw_conn.autocommit = True
            cur = raw_conn.cursor()

            # Get existing values for the subscriptionstatus enum
            cur.execute(
                "SELECT enumlabel FROM pg_enum "
                "JOIN pg_type ON pg_enum.enumtypid = pg_type.oid "
                "WHERE pg_type.typname = 'subscriptionstatus'"
            )
            existing = {row[0] for row in cur.fetchall()}

            needed = ["requires_action"]
            for val in needed:
                if val not in existing:
                    cur.execute(f"ALTER TYPE subscriptionstatus ADD VALUE '{val}'")
                    print(f"[MIGRATE] Added '{val}' to subscriptionstatus enum")

            cur.close()
            raw_conn.close()
        except Exception as e:
            print(f"[MIGRATE] Enum migration skipped: {e}")