"""jsonb_migration_setweight_add_entrez_crdt_ppdat

迁移内容：
1. 19 列 JSON → JSONB
2. 新增 3 列（print_date [PPDAT], create_date [CRDT], entrez_date [EDAT]）
3. 新增 entry_terms 到 global_tags
4. search_tsv 触发器更新（setweight A/B + authors 已为 jsonb 不再需要 CAST）
5. search_tsv 回填全部已有记录
6. 新增搜索索引（B-tree × 4 + GIN × 4）

Revision ID: cb07d6b1df01
Revises: 6b662a8c5235
Create Date: 2026-07-24 23:11:17.262909
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision: str = 'cb07d6b1df01'
down_revision: Union[str, None] = '6b662a8c5235'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

# ═══ search_tsv 新生成公式（setweight 区分权重）═══
# 注意：setweight() 返回 tsvector，tsvector || tsvector 返回 tsvector，
# 所以不需要外层 to_tsvector() 包裹
_TSVEC = """setweight(to_tsvector('english', COALESCE(NEW.title, '')), 'A') ||
    setweight(to_tsvector('english', COALESCE(NEW.abstract, '')), 'B') ||
    setweight(to_tsvector('english',
        COALESCE(
            (SELECT string_agg(
                value->>'family' || ' ' || COALESCE(value->>'affiliation', ''),
                ' ')
             FROM jsonb_array_elements(NEW.authors)),
            '')
    ), 'A')"""

_UPDATE = """setweight(to_tsvector('english', COALESCE(title, '')), 'A') ||
    setweight(to_tsvector('english', COALESCE(abstract, '')), 'B') ||
    setweight(to_tsvector('english',
        COALESCE(
            (SELECT string_agg(
                value->>'family' || ' ' || COALESCE(value->>'affiliation', ''),
                ' ')
             FROM jsonb_array_elements(authors)),
            '')
    ), 'A')"""

# 需改为 JSONB 的列列表
_JSON_TO_JSONB_COLS = [
    "authors", "pub_types", "mesh_headings", "keywords", "grants",
    "full_text_sections", "trial_reg", "study_design", "pico",
    "retraction_details", "license_info", "pre_extracted_data", "rct_detection",
    "negative_result_details", "ai_summary", "chemical_list", "gene_symbols",
    "databank_list", "suppl_mesh_list",
]


def upgrade() -> None:
    # ─── 0. 先删除依赖 authors 列的触发器（否则 alter_column 被阻塞）───
    op.execute("DROP TRIGGER IF EXISTS trg_global_literature_tsv ON global_literature")

    # ─── 1. 新增 3 列 ───
    op.add_column('global_literature', sa.Column('print_date', sa.Date(), nullable=True))
    op.add_column('global_literature', sa.Column('create_date', sa.DateTime(timezone=True), nullable=True))
    op.add_column('global_literature', sa.Column('entrez_date', sa.DateTime(timezone=True), nullable=True))

    # ─── 2. 19 列 JSON → JSONB（单笔 ALTER TABLE，一次全表重写）───
    for col in _JSON_TO_JSONB_COLS:
        op.alter_column('global_literature', col,
                        existing_type=postgresql.JSON(astext_type=sa.Text()),
                        type_=postgresql.JSONB(astext_type=sa.Text()),
                        existing_nullable=True)

    # ─── 3. entry_terms 到 global_tags ───
    op.add_column('global_tags', sa.Column('entry_terms', postgresql.JSONB(astext_type=sa.Text()), nullable=True))

    # ─── 4. 更新 search_tsv 触发器函数（setweight A/B + authors 已是 jsonb 不需 CAST）───
    op.execute(f"""
        CREATE OR REPLACE FUNCTION update_literature_search_tsv()
        RETURNS trigger AS $$
        BEGIN
            NEW.search_tsv := {_TSVEC};
            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """)

    # ─── 4b. 重建触发器 ───
    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()
    """)

    # ─── 5. 回填全部已有记录的 search_tsv ───
    op.execute(f"""
        UPDATE global_literature
        SET search_tsv = {_UPDATE}
        WHERE search_tsv IS NOT NULL;
    """)

    # ─── 6. 新增 B-tree 索引（先删后建，应对部分已存在的索引）───
    op.execute("DROP INDEX IF EXISTS ix_gl_doi")
    op.create_index(op.f('ix_gl_doi'), 'global_literature', ['doi'], unique=False)
    op.execute("DROP INDEX IF EXISTS ix_gl_journal_iso")
    op.create_index(op.f('ix_gl_journal_iso'), 'global_literature', ['journal_iso'], unique=False)
    op.execute("DROP INDEX IF EXISTS ix_gl_language")
    op.create_index(op.f('ix_gl_language'), 'global_literature', ['language'], unique=False)
    op.execute("DROP INDEX IF EXISTS ix_gl_retracted")
    op.create_index(op.f('ix_gl_retracted'), 'global_literature', ['retracted'], unique=False)

    # ─── 7. 新增 GIN 索引（JSONB 列，先删后建）───
    op.execute("DROP INDEX IF EXISTS ix_gl_pub_types_gin")
    op.create_index(op.f('ix_gl_pub_types_gin'), 'global_literature', ['pub_types'],
                    unique=False, postgresql_using='gin')
    op.execute("DROP INDEX IF EXISTS ix_gl_study_design_gin")
    op.create_index(op.f('ix_gl_study_design_gin'), 'global_literature', ['study_design'],
                    unique=False, postgresql_using='gin')
    op.execute("DROP INDEX IF EXISTS ix_gl_grants_gin")
    op.create_index(op.f('ix_gl_grants_gin'), 'global_literature', ['grants'],
                    unique=False, postgresql_using='gin')
    op.execute("DROP INDEX IF EXISTS ix_gl_chemical_list_gin")
    op.create_index(op.f('ix_gl_chemical_list_gin'), 'global_literature', ['chemical_list'],
                    unique=False, postgresql_using='gin')


def downgrade() -> None:
    # ─── 回滚 7. GIN 索引 ───
    op.drop_index(op.f('ix_gl_chemical_list_gin'), table_name='global_literature')
    op.drop_index(op.f('ix_gl_grants_gin'), table_name='global_literature')
    op.drop_index(op.f('ix_gl_study_design_gin'), table_name='global_literature')
    op.drop_index(op.f('ix_gl_pub_types_gin'), table_name='global_literature')

    # ─── 回滚 6. B-tree 索引 ───
    op.drop_index(op.f('ix_gl_retracted'), table_name='global_literature')
    op.drop_index(op.f('ix_gl_language'), table_name='global_literature')
    op.drop_index(op.f('ix_gl_journal_iso'), table_name='global_literature')
    op.drop_index(op.f('ix_gl_doi'), table_name='global_literature')

    # ─── 回滚 0/4. 先删除触发器，恢复旧函数和旧触发器 ───
    op.execute("DROP TRIGGER IF EXISTS trg_global_literature_tsv ON global_literature")
    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' || ' ' || COALESCE(value->>'affiliation', ''),
                        ' ')
                     FROM jsonb_array_elements(CAST(NEW.authors AS jsonb))),
                    ''
                )
            );
            RETURN NEW;
        END;
        $$ LANGUAGE plpgsql;
    """)
    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()
    """)
    op.execute("""
        UPDATE global_literature
        SET search_tsv = to_tsvector('english',
            COALESCE(title, '') || ' ' || COALESCE(abstract, '') || ' ' ||
            COALESCE(
                (SELECT string_agg(
                    value->>'family' || ' ' || COALESCE(value->>'affiliation', ''),
                    ' ')
                 FROM jsonb_array_elements(CAST(authors AS jsonb))),
                ''
            )
        )
        WHERE search_tsv IS NOT NULL;
    """)

    # ─── 回滚 3. entry_terms ───
    op.drop_column('global_tags', 'entry_terms')

    # ─── 回滚 1. 新增 3 列 ───
    op.drop_column('global_literature', 'entrez_date')
    op.drop_column('global_literature', 'create_date')
    op.drop_column('global_literature', 'print_date')

    # ─── 回滚 2. 19 列 JSONB → JSON ───
    for col in reversed(_JSON_TO_JSONB_COLS):
        op.alter_column('global_literature', col,
                        existing_type=postgresql.JSONB(astext_type=sa.Text()),
                        type_=postgresql.JSON(astext_type=sa.Text()),
                        existing_nullable=True)
