"""add authors to search tsvector

Revision ID: 01563ae6211c
Revises: a0a795b041f4
Create Date: 2026-07-10 17:41:02.542729
"""
from typing import Sequence, Union

from alembic import op

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

# authors 列定义为 JSON（不是 JSONB），需显式 CAST
_AUTHORS_TSVEC = """to_tsvector('english',
    COALESCE(NEW.title, '') || ' ' || COALESCE(NEW.abstract, '') || ' ' ||
    COALESCE(
        (SELECT string_agg(value->>'family', ' ')
         FROM jsonb_array_elements(CAST(NEW.authors AS jsonb))),
        ''
    )
)"""

_AUTHORS_UPDATE = """to_tsvector('english',
    COALESCE(title, '') || ' ' || COALESCE(abstract, '') || ' ' ||
    COALESCE(
        (SELECT string_agg(value->>'family', ' ')
         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 := {_AUTHORS_TSVEC};
            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """)

    # 重建触发器，新增 UPDATE OF authors
    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, authors
            ON global_literature
            FOR EACH ROW
            EXECUTE FUNCTION update_literature_search_tsv();
    """)

    # 回填（authors 列是 JSON 类型，需 CAST 后 jsonb_array_elements）
    op.execute(f"""
        UPDATE global_literature
        SET search_tsv = {_AUTHORS_UPDATE}
        WHERE search_tsv IS NULL
           OR (authors IS NOT NULL AND CAST(authors AS jsonb) != '[]'::jsonb);
    """)


def downgrade() -> None:
    # 恢复原状：仅 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 NOT NULL;
    """)
