"""make transforms global

Revision ID: 8a2f6f0802a3
Revises: 6273338ffe42
Create Date: 2025-11-26 16:03:01.167131

"""

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

# revision identifiers, used by Alembic.
revision = "8a2f6f0802a3"
down_revision = "6273338ffe42"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "trace_transforms",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("task_id", sa.String(), nullable=False),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("description", sa.String(), nullable=True),
        sa.Column("definition", postgresql.JSON(astext_type=sa.Text()), nullable=False),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.Column("updated_at", sa.TIMESTAMP(), nullable=False),
        sa.ForeignKeyConstraint(["task_id"], ["tasks.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_trace_transforms_task_id"),
        "trace_transforms",
        ["task_id"],
        unique=False,
    )

    # Migrate existing dataset_transforms data to trace_transforms
    # First, get a connection to execute raw SQL
    connection = op.get_bind()

    # Migrate data: extract from dataset_transforms, transform JSON, insert into trace_transforms
    # Transform the definition JSON:
    # 1. Rename "columns" array to "variables"
    # 2. Within each array element, rename "column_name" field to "variable_name"
    connection.execute(
        sa.text(
            """
        INSERT INTO trace_transforms (id, task_id, name, description, definition, created_at, updated_at)
        SELECT
            gen_random_uuid() as id,
            d.task_id,
            dt.name,
            dt.description,
            json_build_object(
                'variables',
                (
                    SELECT json_agg(
                        json_build_object(
                            'variable_name', col->>'column_name',
                            'span_name', col->>'span_name',
                            'attribute_path', col->>'attribute_path',
                            'fallback', col->>'fallback'
                        )
                    )
                    FROM json_array_elements(dt.definition->'columns') AS col
                )
            ) as definition,
            dt.created_at,
            dt.updated_at
        FROM dataset_transforms dt
        JOIN datasets d ON dt.dataset_id = d.id
    """,
        ),
    )

    # Drop the dataset_transforms table
    op.drop_table("dataset_transforms")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    # Recreate the old dataset_transforms table (empty)
    op.create_table(
        "dataset_transforms",
        sa.Column("id", sa.UUID(), nullable=False),
        sa.Column("dataset_id", sa.UUID(), nullable=False),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("description", sa.String(), nullable=True),
        sa.Column("definition", postgresql.JSON(astext_type=sa.Text()), nullable=False),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.Column("updated_at", sa.TIMESTAMP(), nullable=False),
        sa.ForeignKeyConstraint(["dataset_id"], ["datasets.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "dataset_id",
            "name",
            name="uq_dataset_transforms_dataset_id_name",
        ),
    )
    op.create_index(
        "idx_dataset_transforms_dataset_id",
        "dataset_transforms",
        ["dataset_id"],
        unique=False,
    )

    op.drop_index(op.f("ix_trace_transforms_task_id"), table_name="trace_transforms")
    op.drop_index(op.f("ix_trace_transforms_id"), table_name="trace_transforms")
    op.drop_table("trace_transforms")
    # ### end Alembic commands ###
