"""Give articles and flashcards the same retrieval columns as questions.

Both were searched with `ILIKE '%term%'`, so "yellow newborn" could not find a
jaundice article. They now carry a full-text vector and an embedding, with the
same provenance columns that make a model change visible.

Revision ID: u3a4b5c6d7e8
Revises: t2f3a4b5c697
"""
from alembic import op

revision = "u3a4b5c6d7e8"
down_revision = "t2f3a4b5c697"
branch_labels = None
depends_on = None

TARGETS = {
    # table: expression whose text is indexed, weighted title/name first
    "articles": """
        setweight(to_tsvector('english', coalesce(title, '')), 'A') ||
        setweight(to_tsvector('english', coalesce(summary, '')), 'B') ||
        setweight(to_tsvector('english', coalesce(content, '')), 'C')
    """,
    "flashcards": """
        setweight(to_tsvector('english', coalesce(front, '')), 'A') ||
        setweight(to_tsvector('english', coalesce(back, '')), 'B')
    """,
}


def upgrade():
    for table, expression in TARGETS.items():
        op.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS embedding vector(1024)")
        op.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS embedding_model VARCHAR(120)")
        op.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS embedded_at TIMESTAMP")
        op.execute(f"""
            ALTER TABLE {table} ADD COLUMN IF NOT EXISTS search_vector tsvector
            GENERATED ALWAYS AS ({expression}) STORED
        """)
        op.execute(f"CREATE INDEX IF NOT EXISTS ix_{table}_search_vector ON {table} USING GIN (search_vector)")
        op.execute(f"CREATE INDEX IF NOT EXISTS ix_{table}_embedding_model ON {table}(embedding_model)")


def downgrade():
    for table in TARGETS:
        op.execute(f"DROP INDEX IF EXISTS ix_{table}_embedding_model")
        op.execute(f"DROP INDEX IF EXISTS ix_{table}_search_vector")
        for column in ("search_vector", "embedded_at", "embedding_model", "embedding"):
            op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS {column}")