"""add continuous evals to agentic annotations

Revision ID: 44cd9fd55df9
Revises: 714d70516fae
Create Date: 2025-12-04 11:51:34.719923

"""

import sqlalchemy as sa

from alembic import op

# revision identifiers, used by Alembic.
revision = "44cd9fd55df9"
down_revision = "a1b2c3d4e5f67890"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Drop unique index on trace_id to allow multiple annotations per trace
    op.drop_index(
        op.f("ix_agentic_annotations_trace_id"),
        table_name="agentic_annotations",
    )
    # Recreate as non-unique index
    op.create_index(
        op.f("ix_agentic_annotations_trace_id"),
        "agentic_annotations",
        ["trace_id"],
        unique=False,
    )

    # Add annotation_type as nullable first
    op.add_column(
        "agentic_annotations",
        sa.Column("annotation_type", sa.String(), nullable=True),
    )

    # Set all previously existing annotations to "human"
    op.execute(
        "UPDATE agentic_annotations SET annotation_type = 'human' WHERE annotation_type IS NULL",
    )

    # Now make it non-nullable
    op.alter_column("agentic_annotations", "annotation_type", nullable=False)

    op.add_column(
        "agentic_annotations",
        sa.Column("continuous_eval_id", sa.UUID(), nullable=True),
    )
    op.add_column(
        "agentic_annotations",
        sa.Column("input_variables", sa.JSON(), nullable=True),
    )
    op.add_column("agentic_annotations", sa.Column("cost", sa.Float(), nullable=True))
    op.add_column(
        "agentic_annotations",
        sa.Column("run_status", sa.String(), nullable=True),
    )
    op.alter_column(
        "agentic_annotations",
        "annotation_score",
        existing_type=sa.INTEGER(),
        nullable=True,
    )
    op.create_index(
        op.f("ix_agentic_annotations_continuous_eval_id"),
        "agentic_annotations",
        ["continuous_eval_id"],
        unique=False,
    )
    op.create_foreign_key(
        "fk_agentic_annotations_continuous_eval_id",
        "agentic_annotations",
        "continuous_evals",
        ["continuous_eval_id"],
        ["id"],
        ondelete="CASCADE",
    )

    # Drop existing check constraint and recreate with nullable support
    op.drop_constraint(
        "ck_annotation_score_binary",
        "agentic_annotations",
        type_="check",
    )
    op.create_check_constraint(
        "ck_annotation_score_binary",
        "agentic_annotations",
        "annotation_score IS NULL OR annotation_score IN (0, 1)",
    )

    # Add new check constraints
    op.create_check_constraint(
        "ck_human_requires_score",
        "agentic_annotations",
        "(annotation_type != 'human') OR (annotation_score IS NOT NULL)",
    )
    op.create_check_constraint(
        "ck_continuous_eval_requires_id",
        "agentic_annotations",
        "(annotation_type != 'continuous_eval') OR (continuous_eval_id IS NOT NULL)",
    )
    op.create_check_constraint(
        "ck_continuous_eval_requires_run_status",
        "agentic_annotations",
        "(annotation_type != 'continuous_eval') OR (run_status IS NOT NULL)",
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    # Delete all non-human annotation type rows before downgrade
    op.execute("DELETE FROM agentic_annotations WHERE annotation_type != 'human'")

    # Drop new check constraints
    op.drop_constraint(
        "ck_continuous_eval_requires_run_status",
        "agentic_annotations",
        type_="check",
    )
    op.drop_constraint(
        "ck_continuous_eval_requires_id",
        "agentic_annotations",
        type_="check",
    )
    op.drop_constraint("ck_human_requires_score", "agentic_annotations", type_="check")

    # Restore original check constraint
    op.drop_constraint(
        "ck_annotation_score_binary",
        "agentic_annotations",
        type_="check",
    )
    op.create_check_constraint(
        "ck_annotation_score_binary",
        "agentic_annotations",
        "annotation_score IN (0, 1)",
    )

    op.drop_constraint(
        "fk_agentic_annotations_continuous_eval_id",
        "agentic_annotations",
        type_="foreignkey",
    )
    op.drop_index(
        op.f("ix_agentic_annotations_continuous_eval_id"),
        table_name="agentic_annotations",
    )
    op.alter_column(
        "agentic_annotations",
        "annotation_score",
        existing_type=sa.INTEGER(),
        nullable=False,
    )
    op.drop_column("agentic_annotations", "run_status")
    op.drop_column("agentic_annotations", "cost")
    op.drop_column("agentic_annotations", "input_variables")
    op.drop_column("agentic_annotations", "continuous_eval_id")
    op.drop_column("agentic_annotations", "annotation_type")

    # Recreate unique index on trace_id
    op.drop_index(
        op.f("ix_agentic_annotations_trace_id"),
        table_name="agentic_annotations",
    )
    op.create_index(
        op.f("ix_agentic_annotations_trace_id"),
        "agentic_annotations",
        ["trace_id"],
        unique=True,
    )
    # ### end Alembic commands ###
