"""add transform map to ce

Revision ID: aeea38580473
Revises: 8aba08c2934f
Create Date: 2025-12-16 10:49:49.325132

"""

import json

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

# revision identifiers, used by Alembic.
revision = "aeea38580473"
down_revision = "db04d47f9dd1"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "continuous_evals",
        sa.Column(
            "transform_variable_mapping",
            postgresql.JSON(astext_type=sa.Text()),
            server_default="[]",
            nullable=False,
        ),
    )
    # ### end Alembic commands ###

    # Data migration: populate transform_variable_mapping for existing continuous_evals
    connection = op.get_bind()

    # Get all continuous_evals with their transform and llm_eval info
    continuous_evals = connection.execute(
        sa.text(
            """
            SELECT
                ce.id,
                ce.transform_id,
                ce.task_id,
                ce.llm_eval_name,
                ce.llm_eval_version
            FROM continuous_evals ce
            WHERE ce.transform_variable_mapping::text = '[]'
        """,
        ),
    ).fetchall()

    for ce in continuous_evals:
        ce_id = ce[0]
        transform_id = ce[1]
        task_id = ce[2]
        llm_eval_name = ce[3]
        llm_eval_version = ce[4]

        # Get transform definition to extract variable names
        transform_result = connection.execute(
            sa.text(
                """
                SELECT definition
                FROM trace_transforms
                WHERE id = :transform_id
            """,
            ),
            {"transform_id": transform_id},
        ).fetchone()

        if not transform_result:
            continue

        transform_definition = transform_result[0]

        # Get LLM eval variables
        llm_eval_result = connection.execute(
            sa.text(
                """
                SELECT variables
                FROM llm_evals
                WHERE task_id = :task_id
                  AND name = :name
                  AND version = :version
            """,
            ),
            {"task_id": task_id, "name": llm_eval_name, "version": llm_eval_version},
        ).fetchone()

        if not llm_eval_result:
            continue

        eval_variables = llm_eval_result[0]

        # Extract transform variable names from definition
        transform_variables = transform_definition.get("variables", [])
        transform_var_names = [
            var.get("variable_name")
            for var in transform_variables
            if var.get("variable_name")
        ]

        # Create mapping by matching variables based on name equality
        # Use sets to track which variables have been mapped
        mapped_transform_vars = set()
        mapped_llm_vars = set()
        mapping = []

        for eval_var in eval_variables:
            if (
                eval_var not in mapped_llm_vars
                and eval_var in transform_var_names
                and eval_var not in mapped_transform_vars
            ):
                mapping.append(
                    {"transform_variable": eval_var, "eval_variable": eval_var},
                )
                mapped_transform_vars.add(eval_var)
                mapped_llm_vars.add(eval_var)

        # Update the continuous_eval with the mapping
        if mapping:
            connection.execute(
                sa.text(
                    """
                    UPDATE continuous_evals
                    SET transform_variable_mapping = :mapping
                    WHERE id = :ce_id
                """,
                ),
                {"mapping": json.dumps(mapping), "ce_id": ce_id},
            )


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("continuous_evals", "transform_variable_mapping")
    # ### end Alembic commands ###
