"""add affiliation to search tsvector

Revision ID: e341edea85e2
Revises: 01563ae6211c
Create Date: 2026-07-10 18:01:47.175217
"""
from typing import Sequence, Union

from alembic import op

revision: str = 'e341edea85e2'
down_revision: Union[str, None] = '01563ae6211c'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

# 在已有 family 的基础上，同时提取 affiliation（一次 jsonb_array_elements 扫描）
_TSVEC = """to_tsvector('english',
    COALESCE(NEW.title, '') || ' ' || COALESCE(NEW.abstract, '') || ' ' ||
    COALESCE(
        (SELECT string_agg(
            value->>'family' || ' ' || COALESCE(value->>'affiliation', ''),
            ' ')
         FROM jsonb_array_elements(CAST(NEW.authors AS jsonb))),
        ''
    )
)"""

_UPDATE = """to_tsvector('english',
    COALESCE(title, '') || ' ' || COALESCE(abstract, '') || ' ' ||
    COALESCE(
        (SELECT string_agg(
            value->>'family' || ' ' || COALESCE(value->>'affiliation', ''),
            ' ')
         FROM jsonb_array_elements(CAST(authors AS jsonb))),
        ''
    )
)"""


def upgrade() -> None:
    op.execute(f"""
        CREATE OR REPLACE FUNCTION update_literature_search_tsv()
        RETURNS trigger AS $$
        BEGIN
            NEW.search_tsv := {_TSVEC};
            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """)

    # 触发器 trg_global_literature_tsv 已覆盖 UPDATE OF title, abstract, authors，无需改动

    # 回填全部已有记录（重新计算 tsvector，包含 affiliation）
    op.execute(f"""
        UPDATE global_literature
        SET search_tsv = {_UPDATE}
        WHERE search_tsv IS NOT NULL;
    """)


def downgrade() -> None:
    # 恢复：仅 family，不带 affiliation
    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, '') || ' ' ||
                COALESCE(
                    (SELECT string_agg(value->>'family', ' ')
                     FROM jsonb_array_elements(CAST(NEW.authors AS jsonb))),
                    ''
                )
            );
            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """)

    op.execute("""
        UPDATE global_literature
        SET search_tsv = to_tsvector('english',
            COALESCE(title, '') || ' ' || COALESCE(abstract, '') || ' ' ||
            COALESCE(
                (SELECT string_agg(value->>'family', ' ')
                 FROM jsonb_array_elements(CAST(authors AS jsonb))),
                ''
            )
        )
        WHERE search_tsv IS NOT NULL;
    """)
