"""add search_tsv tsvector column + GIN index

Revision ID: e8f9a0b1c2d3
Revises: bb7bee84037c
Create Date: 2026-07-08 08:20:00.000000
"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import TSVECTOR

revision: str = 'e8f9a0b1c2d3'
down_revision: Union[str, None] = 'bb7bee84037c'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # tsvector 列
    op.add_column(
        'global_literature',
        sa.Column('search_tsv', TSVECTOR(), nullable=True),
    )

    # GIN 索引
    op.create_index(
        'ix_gl_search_tsv',
        'global_literature',
        ['search_tsv'],
        postgresql_using='gin',
    )

    # 触发器函数：INSERT / UPDATE title/abstract 时自动维护
    op.execute("""
        CREATE OR REPLACE FUNCTION update_literature_search_tsv()
        RETURNS trigger AS $$
        BEGIN
            NEW.search_tsv := to_tsvector('english',
                COALESCE(NEW.title, '') || ' ' || COALESCE(NEW.abstract, '')
            );
            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """)

    op.execute("DROP TRIGGER IF EXISTS trg_global_literature_tsv ON global_literature;")
    op.execute("""
        CREATE TRIGGER trg_global_literature_tsv
            BEFORE INSERT OR UPDATE OF title, abstract
            ON global_literature
            FOR EACH ROW
            EXECUTE FUNCTION update_literature_search_tsv();
    """)

    # 回填已有记录
    op.execute("""
        UPDATE global_literature
        SET search_tsv = to_tsvector('english',
            COALESCE(title, '') || ' ' || COALESCE(abstract, '')
        )
        WHERE search_tsv IS NULL;
    """)


def downgrade() -> None:
    op.execute("DROP TRIGGER IF EXISTS trg_global_literature_tsv ON global_literature;")
    op.execute("DROP FUNCTION IF EXISTS update_literature_search_tsv();")
    op.drop_index('ix_gl_search_tsv', table_name='global_literature')
    op.drop_column('global_literature', 'search_tsv')
