app/core/database.py

from collections.abc import Generator

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

from app.core.config import Settings, get_settings


settings = get_settings()


def _sqlite_url(database_url: str) -> bool:
    return database_url.startswith("sqlite")


def _pool_setting(value: int, *, minimum: int) -> int:
    return max(minimum, int(value))


def build_engine_options(active_settings: Settings | None = None) -> dict[str, object]:
    active_settings = active_settings or settings
    options: dict[str, object] = {
        "pool_pre_ping": True,
    }
    if _sqlite_url(active_settings.database_url):
        options["connect_args"] = {"check_same_thread": False}
        return options

    # Supabase's session pooler commonly caps a project at 15 connections.
    # SQLAlchemy's default QueuePool can consume that entire allowance from one
    # process, so keep the app's per-process footprint intentionally small.
    options.update(
        {
            "connect_args": {},
            "pool_size": _pool_setting(active_settings.database_pool_size, minimum=1),
            "max_overflow": _pool_setting(active_settings.database_max_overflow, minimum=0),
            "pool_timeout": _pool_setting(active_settings.database_pool_timeout_seconds, minimum=1),
            "pool_recycle": _pool_setting(active_settings.database_pool_recycle_seconds, minimum=60),
        }
    )
    return options


def _startup_safety_checks() -> None:
    """Prod hardening + Supabase/Postgres best practices (inspired by audit + existing main.py checks).
    Called early from lifespan / init paths.
    """
    from app.core.config import get_settings as _get_settings
    s = _get_settings()
    is_prod = (s.environment or "").lower() == "production"

    if is_prod and _sqlite_url(s.database_url):
        raise RuntimeError(
            "DATABASE_URL is SQLite in production. Use PostgreSQL (Supabase) connection string."
        )

    # Log pool config for observability (Supabase pooler sensitive)
    if not _sqlite_url(s.database_url):
        import logging
        logging.getLogger("docdoe.db").info(
            "DB pool config: size=%s overflow=%s recycle=%ss timeout=%ss (Supabase pooler friendly)",
            s.database_pool_size, s.database_max_overflow, s.database_pool_recycle_seconds, s.database_pool_timeout_seconds,
        )

engine = create_engine(
    settings.database_url,
    **build_engine_options(),
)

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


class Base(DeclarativeBase):
    pass


def get_db() -> Generator[Session, None, None]:
    db = SessionLocal()
    try:
        yield db
    finally:
        # Roll back any pending/aborted transaction so a poisoned connection is
        # cleaned before returning to the pool (prevents InFailedSqlTransaction
        # leaking into the next request that reuses this pooled connection).
        try:
            db.rollback()
        except Exception:  # noqa: BLE001
            pass
        db.close()


def init_db() -> None:
    # Import models so SQLAlchemy registers all tables before create_all.
    from app.models import (  # noqa: F401
        chat_session,
        document,
        document_chunk,
        flashcard,
        generation_cache,
        generation,
        previous_paper,
        previous_question,
        provider_usage_log,
        quiz,
        study_profile,
        student_workspace,
        tuition_profile,
        class_session_progress,
        phase3_activity,
        password_reset_token,
        learn_anything_roadmap,
        learning_state,
        support_submission,
        syllabus_item,
        chapter_pattern,
        telemetry,
        user,
        user_plan,
        video_render_job,
        weak_topic,
        job,
        user_usage_monthly,
    )
    from app.models.user import User

    # Supabase/Postgres best practices note (adapted from audit):
    # - Prefer proper migrations (alembic or supabase migration new + db push) over repeated create_all in prod.
    # - create_all is convenient for dev/SQLite but can drift; the _ensure_*_columns helpers below act as lightweight "on-startup migrations".
    # - For Supabase: use connection pooler, keep pool small (see build_engine_options), SSL, and monitor limits.
    # - RLS is not enforced here (backend uses service-role via DATABASE_URL); if exposing tables via Data API/PostgREST in future, enable RLS + policies using auth.uid().
    Base.metadata.create_all(bind=engine)
    _startup_safety_checks()  # Supabase/Postgres prod guards + pool notes (from best-practices audit)
    if _ensure_document_columns():
        _backfill_legacy_document_material_types()
    _ensure_document_education_columns()
    _ensure_document_chunk_columns()
    _ensure_previous_paper_columns()
    _ensure_ai_result_columns()
    _ensure_chat_session_columns()
    _ensure_chat_message_columns()
    _ensure_video_render_job_columns()
    _ensure_generation_columns()
    _ensure_user_columns()
    _ensure_user_plan_columns()
    _ensure_subscription_columns()
    _ensure_previous_question_t2_columns()
    _ensure_student_profile_exam_date_nullable()
    _ensure_quiz_attempt_columns()

    with SessionLocal() as db:
        demo_user = db.get(User, "usr_demo_student")
        if demo_user is None:
            db.add(
                User(
                    id="usr_demo_student",
                    name="Demo Student",
                    email="student@example.com",
                    role="student",
                    class_level="Plus Two",
                    syllabus="Kerala HSE",
                    preferred_language="Malayalam + English",
                )
            )
            db.commit()


def _ensure_document_columns() -> bool:
    inspector = inspect(engine)
    if "documents" not in inspector.get_table_names():
        return False

    document_columns = {column["name"] for column in inspector.get_columns("documents")}
    material_type_added = False
    with engine.begin() as connection:
        if "extraction_error" not in document_columns:
            connection.execute(text("ALTER TABLE documents ADD COLUMN extraction_error TEXT"))
        if "chunk_count" not in document_columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN chunk_count INTEGER NOT NULL DEFAULT 0"),
            )
        if "source_type" not in document_columns:
            connection.execute(
                text(
                    "ALTER TABLE documents "
                    "ADD COLUMN source_type TEXT NOT NULL DEFAULT 'pdf'",
                ),
            )
        if "material_type" not in document_columns:
            connection.execute(
                text(
                    "ALTER TABLE documents "
                    "ADD COLUMN material_type TEXT NOT NULL DEFAULT 'unknown'",
                ),
            )
            material_type_added = True
        if "updated_at" not in document_columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN updated_at TIMESTAMP"),
            )
            connection.execute(
                text("UPDATE documents SET updated_at = created_at WHERE updated_at IS NULL"),
            )
        if "extracted_text_length" not in document_columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN extracted_text_length INTEGER"),
            )
        if "processing_started_at" not in document_columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN processing_started_at TIMESTAMP"),
            )
        if "processing_completed_at" not in document_columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN processing_completed_at TIMESTAMP"),
            )
    return material_type_added


def _backfill_legacy_document_material_types() -> None:
    from app.models.document import Document
    from app.services.source_classifier import classify_material_type

    with SessionLocal() as db:
        documents = db.scalars(
            select(Document).where(Document.material_type == "unknown"),
        ).all()
        updated = False
        for document in documents:
            inferred = classify_material_type(document.file_name, document.extracted_text)
            if inferred != "unknown":
                document.material_type = inferred
                updated = True
        if updated:
            db.commit()


def _ensure_document_education_columns() -> None:
    inspector = inspect(engine)
    if "documents" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("documents")}
    with engine.begin() as connection:
        if "education_extraction_status" not in columns:
            connection.execute(
                text(
                    "ALTER TABLE documents "
                    "ADD COLUMN education_extraction_status TEXT NOT NULL DEFAULT 'uploaded'"
                ),
            )
        if "education_extraction_error" not in columns:
            connection.execute(text("ALTER TABLE documents ADD COLUMN education_extraction_error TEXT"))
        if "education_warnings_json" not in columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN education_warnings_json JSON DEFAULT '[]'"),
            )
        if "syllabus_items_count" not in columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN syllabus_items_count INTEGER NOT NULL DEFAULT 0"),
            )
        if "pyq_questions_count" not in columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN pyq_questions_count INTEGER NOT NULL DEFAULT 0"),
            )
        if "pyq_years_json" not in columns:
            connection.execute(
                text("ALTER TABLE documents ADD COLUMN pyq_years_json JSON DEFAULT '[]'"),
            )


def _ensure_ai_result_columns() -> None:
    inspector = inspect(engine)
    tables = set(inspector.get_table_names())
    targets = {
        "quizzes": "questions_json",
        "flashcard_sets": "cards_json",
    }
    with engine.begin() as connection:
        for table_name in targets:
            if table_name not in tables:
                continue
            columns = {column["name"] for column in inspector.get_columns(table_name)}
            if "model_used" not in columns:
                connection.execute(
                    text(
                        f"ALTER TABLE {table_name} "
                        "ADD COLUMN model_used TEXT NOT NULL DEFAULT 'mock-exam-tutor-v1'",
                    ),
                )


def _ensure_chat_session_columns() -> None:
    inspector = inspect(engine)
    if "chat_sessions" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("chat_sessions")}
    if "context_data" not in columns:
        with engine.begin() as connection:
            connection.execute(
                text("ALTER TABLE chat_sessions ADD COLUMN context_data JSON NOT NULL DEFAULT '{}'")
            )


def _ensure_chat_message_columns() -> None:
    inspector = inspect(engine)
    if "chat_messages" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("chat_messages")}
    with engine.begin() as connection:
        if "evidence_label" not in columns:
            connection.execute(text("ALTER TABLE chat_messages ADD COLUMN evidence_label TEXT"))
        if "web_sources" not in columns:
            connection.execute(
                text("ALTER TABLE chat_messages ADD COLUMN web_sources JSON NOT NULL DEFAULT '[]'"),
            )
        if "client_turn_id" not in columns:
            connection.execute(
                text("ALTER TABLE chat_messages ADD COLUMN client_turn_id TEXT"),
            )
        connection.execute(
            text(
                "CREATE UNIQUE INDEX IF NOT EXISTS uq_chat_messages_session_turn_role "
                "ON chat_messages(session_id, client_turn_id, role) "
                "WHERE client_turn_id IS NOT NULL",
            ),
        )


def _ensure_previous_paper_columns() -> None:
    inspector = inspect(engine)
    if "previous_papers" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("previous_papers")}
    with engine.begin() as connection:
        if "file_name" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN file_name TEXT"))
        if "file_type" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN file_type TEXT"))
        if "status" not in columns:
            connection.execute(
                text("ALTER TABLE previous_papers ADD COLUMN status TEXT NOT NULL DEFAULT 'ready'"),
            )
        if "extracted_text" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN extracted_text TEXT"))
        if "extraction_error" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN extraction_error TEXT"))
        if "board" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN board TEXT"))
        if "class_level" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN class_level TEXT"))
        if "source_url" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN source_url TEXT"))
        if "source_domain" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN source_domain TEXT"))
        if "source_title" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN source_title TEXT"))
        if "retrieved_at" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN retrieved_at TIMESTAMP"))
        if "file_hash" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN file_hash TEXT"))
        if "verification_status" not in columns:
            connection.execute(
                text(
                    "ALTER TABLE previous_papers "
                    "ADD COLUMN verification_status TEXT NOT NULL DEFAULT 'verified'",
                ),
            )
        if "confidence_score" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN confidence_score FLOAT"))
        if "official_source" not in columns:
            connection.execute(
                text(
                    "ALTER TABLE previous_papers "
                    "ADD COLUMN official_source BOOLEAN NOT NULL DEFAULT FALSE",
                ),
            )
        if "notes" not in columns:
            connection.execute(text("ALTER TABLE previous_papers ADD COLUMN notes TEXT"))


def _ensure_video_render_job_columns() -> None:
    inspector = inspect(engine)
    if "video_render_jobs" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("video_render_jobs")}
    with engine.begin() as connection:
        if "output_object_key" not in columns:
            connection.execute(text("ALTER TABLE video_render_jobs ADD COLUMN output_object_key TEXT"))
        if "public_url" not in columns:
            connection.execute(text("ALTER TABLE video_render_jobs ADD COLUMN public_url TEXT"))
        if "storage_provider" not in columns:
            connection.execute(text("ALTER TABLE video_render_jobs ADD COLUMN storage_provider TEXT"))
        if "source_document_id" not in columns:
            connection.execute(text("ALTER TABLE video_render_jobs ADD COLUMN source_document_id TEXT"))
        if "evidence_label" not in columns:
            connection.execute(text("ALTER TABLE video_render_jobs ADD COLUMN evidence_label TEXT"))
        if "target_duration_seconds" not in columns:
            connection.execute(text("ALTER TABLE video_render_jobs ADD COLUMN target_duration_seconds FLOAT"))
        if "audio_duration_seconds" not in columns:
            connection.execute(text("ALTER TABLE video_render_jobs ADD COLUMN audio_duration_seconds FLOAT"))
        if "render_duration_seconds" not in columns:
            connection.execute(text("ALTER TABLE video_render_jobs ADD COLUMN render_duration_seconds FLOAT"))
        if "scene_audio_statuses_json" not in columns:
            connection.execute(
                text("ALTER TABLE video_render_jobs ADD COLUMN scene_audio_statuses_json JSON"),
            )


def _ensure_student_profile_exam_date_nullable() -> None:
    """Make student_profiles.exam_date nullable.

    Exam date is optional at onboarding ("I don't know my exam date yet") — a
    student must never be trapped. The column was originally NOT NULL, so relax
    it idempotently on an existing table (create_all never alters columns).

    The formal Postgres migration is
    ``supabase/migrations/20260718000000_exam_date_nullable.sql`` (applied via
    ``supabase db push``). This startup shim is kept because (a) SQLite dev/CI
    databases are created by create_all and Supabase migrations never run there,
    and (b) it guarantees the running app is consistent even if a deploy reaches
    a Postgres instance before the migration has been pushed. On Postgres it is
    the same idempotent ``DROP NOT NULL`` as the migration.
    """
    inspector = inspect(engine)
    if "student_profiles" not in inspector.get_table_names():
        return
    exam_col = next(
        (c for c in inspector.get_columns("student_profiles") if c["name"] == "exam_date"),
        None,
    )
    if exam_col is None or exam_col.get("nullable", True):
        return  # already nullable (or absent) — nothing to do

    if engine.dialect.name == "postgresql":
        try:
            with engine.begin() as connection:
                connection.execute(
                    text("ALTER TABLE student_profiles ALTER COLUMN exam_date DROP NOT NULL")
                )
        except Exception:
            pass
        return

    if engine.dialect.name == "sqlite":
        # SQLite cannot ALTER a column's nullability in place; rebuild the table
        # with the same columns but a nullable exam_date, preserving all rows.
        try:
            with engine.begin() as connection:
                cols = inspector.get_columns("student_profiles")
                col_names = ", ".join(f'"{c["name"]}"' for c in cols)
                col_defs = []
                for c in cols:
                    coltype = c["type"].compile(dialect=engine.dialect)
                    nn = "" if c["name"] == "exam_date" else (" NOT NULL" if not c.get("nullable", True) else "")
                    pk = " PRIMARY KEY" if c.get("primary_key") else ""
                    default = c.get("default")
                    dflt = f" DEFAULT {default}" if default is not None else ""
                    col_defs.append(f'"{c["name"]}" {coltype}{pk}{dflt}{nn}')
                connection.execute(text("PRAGMA foreign_keys=OFF"))
                connection.execute(text("ALTER TABLE student_profiles RENAME TO student_profiles_old"))
                connection.execute(text(f'CREATE TABLE student_profiles ({", ".join(col_defs)})'))
                connection.execute(
                    text(f"INSERT INTO student_profiles ({col_names}) SELECT {col_names} FROM student_profiles_old")
                )
                connection.execute(text("DROP TABLE student_profiles_old"))
                connection.execute(text("PRAGMA foreign_keys=ON"))
        except Exception:
            pass


def _ensure_quiz_attempt_columns() -> None:
    """Backfill assessment idempotency on existing development databases."""

    inspector = inspect(engine)
    if "quiz_attempts" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("quiz_attempts")}
    with engine.begin() as connection:
        if "client_attempt_id" not in columns:
            connection.execute(
                text("ALTER TABLE quiz_attempts ADD COLUMN client_attempt_id VARCHAR(180)")
            )
        connection.execute(
            text(
                "CREATE UNIQUE INDEX IF NOT EXISTS uq_quiz_attempts_user_client "
                "ON quiz_attempts (user_id, client_attempt_id)"
            )
        )


def _ensure_user_columns() -> None:
    inspector = inspect(engine)
    if "users" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("users")}
    with engine.begin() as connection:
        if "password_hash" not in columns:
            connection.execute(text("ALTER TABLE users ADD COLUMN password_hash TEXT"))
        if "auth_version" not in columns:
            connection.execute(
                text("ALTER TABLE users ADD COLUMN auth_version INTEGER NOT NULL DEFAULT 1")
            )


def _ensure_document_chunk_columns() -> None:
    inspector = inspect(engine)
    if "document_chunks" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("document_chunks")}
    with engine.begin() as connection:
        if "embedding" not in columns:
            connection.execute(text("ALTER TABLE document_chunks ADD COLUMN embedding TEXT"))


def _ensure_user_plan_columns() -> None:
    inspector = inspect(engine)
    if "user_plans" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("user_plans")}
    with engine.begin() as connection:
        if "period_start" not in columns:
            connection.execute(text("ALTER TABLE user_plans ADD COLUMN period_start TIMESTAMP"))
            connection.execute(
                text("UPDATE user_plans SET period_start = created_at WHERE period_start IS NULL")
            )


def _ensure_subscription_columns() -> None:
    """Backfill Stripe lifecycle columns on existing local/hosted databases.

    Explicit SQL migrations remain the production source of truth. This guard
    keeps SQLite development databases usable when they predate that migration.
    """

    inspector = inspect(engine)
    if "subscriptions" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("subscriptions")}
    with engine.begin() as connection:
        if "provider_price_id" not in columns:
            connection.execute(text("ALTER TABLE subscriptions ADD COLUMN provider_price_id TEXT"))
        if "cancel_at_period_end" not in columns:
            connection.execute(
                text(
                    "ALTER TABLE subscriptions "
                    "ADD COLUMN cancel_at_period_end BOOLEAN NOT NULL DEFAULT FALSE"
                )
            )
        connection.execute(
            text(
                "CREATE UNIQUE INDEX IF NOT EXISTS ix_subscriptions_provider_customer "
                "ON subscriptions (provider_customer_id)"
            )
        )
        connection.execute(
            text(
                "CREATE UNIQUE INDEX IF NOT EXISTS ix_subscriptions_provider_subscription "
                "ON subscriptions (provider_subscription_id)"
            )
        )


def _ensure_generation_columns() -> None:
    inspector = inspect(engine)
    if "generations" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("generations")}
    with engine.begin() as connection:
        if "provider_used" not in columns:
            connection.execute(text("ALTER TABLE generations ADD COLUMN provider_used TEXT"))
        if "generation_time_ms" not in columns:
            connection.execute(text("ALTER TABLE generations ADD COLUMN generation_time_ms INTEGER"))
        if "is_mock_output" not in columns:
            connection.execute(text("ALTER TABLE generations ADD COLUMN is_mock_output BOOLEAN"))
        if "validation_status" not in columns:
            connection.execute(text("ALTER TABLE generations ADD COLUMN validation_status TEXT"))
        if "validation_error" not in columns:
            connection.execute(text("ALTER TABLE generations ADD COLUMN validation_error TEXT"))


def _ensure_previous_question_t2_columns() -> None:
    """T2: Add PYQ extraction fields to previous_questions table."""
    inspector = inspect(engine)
    if "previous_questions" not in inspector.get_table_names():
        return

    columns = {column["name"] for column in inspector.get_columns("previous_questions")}
    with engine.begin() as connection:
        if "answer_type" not in columns:
            connection.execute(text("ALTER TABLE previous_questions ADD COLUMN answer_type TEXT"))
        if "formula_needed" not in columns:
            connection.execute(
                text("ALTER TABLE previous_questions ADD COLUMN formula_needed BOOLEAN NOT NULL DEFAULT FALSE"),
            )
        if "diagram_needed" not in columns:
            connection.execute(
                text("ALTER TABLE previous_questions ADD COLUMN diagram_needed BOOLEAN NOT NULL DEFAULT FALSE"),
            )
        if "extracted_answer_if_available" not in columns:
            connection.execute(text("ALTER TABLE previous_questions ADD COLUMN extracted_answer_if_available TEXT"))
        if "confidence" not in columns:
            connection.execute(
                text("ALTER TABLE previous_questions ADD COLUMN confidence FLOAT NOT NULL DEFAULT 0.0"),
            )
        if "source_origin" not in columns:
            connection.execute(
                text("ALTER TABLE previous_questions ADD COLUMN source_origin TEXT NOT NULL DEFAULT 'user_uploaded'"),
            )