"""refactor rag tags

Revision ID: 52c0a0da3ef1
Revises: d31e27429dca
Create Date: 2025-11-06 16:04:19.781360

"""

import json
import uuid

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

# revision identifiers, used by Alembic.
revision = "52c0a0da3ef1"
down_revision = "d31e27429dca"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Create the new table first
    op.create_table(
        "rag_search_setting_tags",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("setting_configuration_id", sa.UUID(), nullable=False),
        sa.Column("version_number", sa.Integer(), nullable=False),
        sa.Column("tag", sa.String(), nullable=False),
        sa.ForeignKeyConstraint(
            ["setting_configuration_id", "version_number"],
            [
                "rag_search_setting_configuration_versions.setting_configuration_id",
                "rag_search_setting_configuration_versions.version_number",
            ],
        ),
        sa.ForeignKeyConstraint(
            ["setting_configuration_id"],
            ["rag_search_setting_configurations.id"],
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "setting_configuration_id",
            "tag",
            name="uix_unique_setting_configuration_id_tag",
        ),
    )
    op.create_index(
        op.f("ix_rag_search_setting_tags_setting_configuration_id"),
        "rag_search_setting_tags",
        ["setting_configuration_id"],
        unique=False,
    )
    op.create_index(
        "rag_search_setting_configurations_versions_tags",
        "rag_search_setting_tags",
        ["setting_configuration_id", "version_number"],
        unique=False,
    )

    # Migrate data from old JSON columns to new table
    # Get a connection to execute raw SQL
    connection = op.get_bind()

    # Migrate tags from rag_search_setting_configuration_versions.tags
    # The tags column is a JSON array of strings
    # Order by setting_configuration_id and version_number to process in order
    versions_result = connection.execute(
        sa.text(
            """
        SELECT setting_configuration_id, version_number, tags
        FROM rag_search_setting_configuration_versions
        WHERE tags IS NOT NULL
        ORDER BY setting_configuration_id, version_number
    """,
        ),
    )

    tag_rows = []
    seen_tags = {}  # Track (setting_configuration_id, tag) to enforce uniqueness

    for row in versions_result:
        setting_configuration_id = row[0]
        version_number = row[1]
        tags_json = row[2]
        # Parse JSON array - handle both string and already-parsed JSON
        if isinstance(tags_json, str):
            tags_list = json.loads(tags_json)
        else:
            tags_list = tags_json

        if tags_list and isinstance(tags_list, list):
            for tag_string in tags_list:
                if isinstance(tag_string, str):
                    tag_key = (str(setting_configuration_id), tag_string)
                    # Only add if we haven't seen this (setting_configuration_id, tag) combination
                    # The unique constraint requires one tag record per (setting_configuration_id, tag)
                    # Use the first version_number where the tag appears
                    if tag_key not in seen_tags:
                        seen_tags[tag_key] = True
                        tag_rows.append(
                            {
                                "id": uuid.uuid4(),
                                "setting_configuration_id": setting_configuration_id,
                                "version_number": version_number,
                                "tag": tag_string,
                            },
                        )

    # Insert all tag rows using executemany for efficiency
    if tag_rows:
        insert_stmt = sa.text(
            """
            INSERT INTO rag_search_setting_tags (id, setting_configuration_id, version_number, tag)
            VALUES (:id, :setting_configuration_id, :version_number, :tag)
        """,
        )
        connection.execute(insert_stmt, tag_rows)

    # Now drop the old columns
    op.drop_column("rag_search_setting_configuration_versions", "tags")
    op.drop_column("rag_search_setting_configurations", "all_possible_tags")


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "rag_search_setting_configurations",
        sa.Column(
            "all_possible_tags",
            postgresql.JSON(astext_type=sa.Text()),
            autoincrement=False,
            nullable=True,
        ),
    )
    op.add_column(
        "rag_search_setting_configuration_versions",
        sa.Column(
            "tags",
            postgresql.JSON(astext_type=sa.Text()),
            autoincrement=False,
            nullable=True,
        ),
    )
    op.drop_index(
        "rag_search_setting_configurations_versions_tags",
        table_name="rag_search_setting_tags",
    )
    op.drop_index(
        op.f("ix_rag_search_setting_tags_setting_configuration_id"),
        table_name="rag_search_setting_tags",
    )
    op.drop_table("rag_search_setting_tags")
    # ### end Alembic commands ###
