"""migrate_tags_array_to_jsonb

Revision ID: 3e29c99b24ae
Revises: 2a19eed6a4b4
Create Date: 2026-01-22 14:32:39.004739

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = '3e29c99b24ae'
down_revision: Union[str, None] = '2a19eed6a4b4'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Migrate tags from ARRAY(String) to JSONB with metadata.

    New JSONB structure per spec:
    [
        {
            "tag_id": "uuid",
            "name": "req:functional",
            "source": "ai",
            "confidence": 0.95,
            "status": "verified",
            "history": [...]
        }
    ]
    """
    from sqlalchemy.dialects.postgresql import JSONB

    # Step 1: Rename old column
    op.alter_column(
        'extracted_records',
        'tags',
        new_column_name='tags_old'
    )

    # Step 2: Drop old GIN index (it's on tags_old now)
    op.drop_index('idx_extracted_records_tags', table_name='extracted_records')

    # Step 3: Create new tags column as JSONB
    op.add_column(
        'extracted_records',
        sa.Column('tags', JSONB, nullable=True, server_default='[]')
    )

    # Step 4: Migrate data - convert ARRAY(String) to JSONB structure
    # Use raw SQL for complex data transformation
    op.execute("""
        UPDATE extracted_records
        SET tags = (
            SELECT jsonb_agg(
                jsonb_build_object(
                    'tag_id', gen_random_uuid()::text,
                    'name', tag_name,
                    'source', 'ai',
                    'confidence', 1.0,
                    'status', 'verified',
                    'history', jsonb_build_array(
                        jsonb_build_object(
                            'action', 'migrated',
                            'by', 'system',
                            'ts', extract(epoch from now())::bigint
                        )
                    )
                )
            )
            FROM unnest(tags_old) AS tag_name
        )
        WHERE tags_old IS NOT NULL AND array_length(tags_old, 1) > 0
    """)

    # Step 5: Drop old column
    op.drop_column('extracted_records', 'tags_old')

    # Step 6: Create new GIN index for JSONB
    op.create_index(
        'idx_extracted_records_tags',
        'extracted_records',
        ['tags'],
        unique=False,
        postgresql_using='gin'
    )


def downgrade() -> None:
    """Revert JSONB tags back to ARRAY(String)."""
    from sqlalchemy.dialects.postgresql import ARRAY

    # Step 1: Drop GIN index on JSONB tags
    op.drop_index('idx_extracted_records_tags', table_name='extracted_records')

    # Step 2: Rename current tags to tags_new
    op.alter_column(
        'extracted_records',
        'tags',
        new_column_name='tags_new'
    )

    # Step 3: Create tags column as ARRAY(String)
    op.add_column(
        'extracted_records',
        sa.Column('tags', ARRAY(sa.String()), nullable=True, server_default='{}')
    )

    # Step 4: Migrate data back - extract 'name' field from JSONB
    op.execute("""
        UPDATE extracted_records
        SET tags = ARRAY(
            SELECT tag->>'name'
            FROM jsonb_array_elements(tags_new) AS tag
        )
        WHERE tags_new IS NOT NULL AND jsonb_array_length(tags_new) > 0
    """)

    # Step 5: Drop new JSONB column
    op.drop_column('extracted_records', 'tags_new')

    # Step 6: Recreate GIN index for ARRAY
    op.create_index(
        'idx_extracted_records_tags',
        'extracted_records',
        ['tags'],
        unique=False,
        postgresql_using='gin'
    )
