"""prompt experiments

Revision ID: 0381fe93ad77
Revises: 5d30b8ab358b
Create Date: 2025-11-15 19:32:26.838827

"""

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = "0381fe93ad77"
down_revision = "5d30b8ab358b"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "prompt_experiments",
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("task_id", sa.String(), nullable=False),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("description", sa.Text(), nullable=True),
        sa.Column("status", sa.String(), nullable=False),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.Column("finished_at", sa.TIMESTAMP(), nullable=True),
        sa.Column("prompt_name", sa.String(), nullable=False),
        sa.Column(
            "prompt_versions", postgresql.JSON(astext_type=sa.Text()), nullable=False
        ),
        sa.Column("dataset_id", sa.String(), nullable=False),
        sa.Column("dataset_version", sa.Integer(), nullable=False),
        sa.Column(
            "prompt_variable_mapping",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column(
            "eval_configs", postgresql.JSON(astext_type=sa.Text()), nullable=False
        ),
        sa.Column("total_rows", sa.Integer(), nullable=False),
        sa.Column("completed_rows", sa.Integer(), nullable=False),
        sa.Column("failed_rows", sa.Integer(), nullable=False),
        sa.Column("total_cost", sa.String(), nullable=True),
        sa.Column(
            "summary_results", postgresql.JSON(astext_type=sa.Text()), nullable=True
        ),
        sa.ForeignKeyConstraint(["task_id"], ["tasks.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_prompt_experiments_created_at"),
        "prompt_experiments",
        ["created_at"],
        unique=False,
    )
    op.create_index(
        op.f("ix_prompt_experiments_id"), "prompt_experiments", ["id"], unique=False
    )
    op.create_index(
        op.f("ix_prompt_experiments_prompt_name"),
        "prompt_experiments",
        ["prompt_name"],
        unique=False,
    )
    op.create_index(
        op.f("ix_prompt_experiments_status"),
        "prompt_experiments",
        ["status"],
        unique=False,
    )
    op.create_index(
        op.f("ix_prompt_experiments_task_id"),
        "prompt_experiments",
        ["task_id"],
        unique=False,
    )
    op.create_table(
        "prompt_experiment_test_cases",
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("experiment_id", sa.String(), nullable=False),
        sa.Column("status", sa.String(), nullable=False),
        sa.Column("dataset_row_id", sa.String(), nullable=False),
        sa.Column(
            "prompt_input_variables",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column("total_cost", sa.String(), nullable=True),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.Column("updated_at", sa.TIMESTAMP(), nullable=False),
        sa.ForeignKeyConstraint(
            ["experiment_id"], ["prompt_experiments.id"], ondelete="CASCADE"
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_prompt_experiment_test_cases_dataset_row_id"),
        "prompt_experiment_test_cases",
        ["dataset_row_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_prompt_experiment_test_cases_experiment_id"),
        "prompt_experiment_test_cases",
        ["experiment_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_prompt_experiment_test_cases_id"),
        "prompt_experiment_test_cases",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_prompt_experiment_test_cases_status"),
        "prompt_experiment_test_cases",
        ["status"],
        unique=False,
    )
    op.create_table(
        "prompt_experiment_test_case_prompt_results",
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("test_case_id", sa.String(), nullable=False),
        sa.Column("name", sa.String(), nullable=False),
        sa.Column("version", sa.Integer(), nullable=False),
        sa.Column("rendered_prompt", sa.Text(), nullable=False),
        sa.Column("output_content", sa.Text(), nullable=True),
        sa.Column(
            "output_tool_calls", postgresql.JSON(astext_type=sa.Text()), nullable=True
        ),
        sa.Column("output_cost", sa.String(), nullable=True),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.ForeignKeyConstraint(
            ["test_case_id"], ["prompt_experiment_test_cases.id"], ondelete="CASCADE"
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_prompt_experiment_test_case_prompt_results_id"),
        "prompt_experiment_test_case_prompt_results",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_prompt_experiment_test_case_prompt_results_test_case_id"),
        "prompt_experiment_test_case_prompt_results",
        ["test_case_id"],
        unique=False,
    )
    op.create_table(
        "prompt_experiment_test_case_prompt_result_eval_scores",
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("prompt_result_id", sa.String(), nullable=False),
        sa.Column("eval_name", sa.String(), nullable=False),
        sa.Column("eval_version", sa.Integer(), nullable=False),
        sa.Column(
            "eval_input_variables",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column("eval_result_score", sa.Float(), nullable=True),
        sa.Column("eval_result_explanation", sa.Text(), nullable=True),
        sa.Column("eval_result_cost", sa.String(), nullable=True),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.ForeignKeyConstraint(
            ["prompt_result_id"],
            ["prompt_experiment_test_case_prompt_results.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_prompt_experiment_test_case_prompt_result_eval_scores_id"),
        "prompt_experiment_test_case_prompt_result_eval_scores",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f(
            "ix_prompt_experiment_test_case_prompt_result_eval_scores_prompt_result_id"
        ),
        "prompt_experiment_test_case_prompt_result_eval_scores",
        ["prompt_result_id"],
        unique=False,
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(
        op.f(
            "ix_prompt_experiment_test_case_prompt_result_eval_scores_prompt_result_id"
        ),
        table_name="prompt_experiment_test_case_prompt_result_eval_scores",
    )
    op.drop_index(
        op.f("ix_prompt_experiment_test_case_prompt_result_eval_scores_id"),
        table_name="prompt_experiment_test_case_prompt_result_eval_scores",
    )
    op.drop_table("prompt_experiment_test_case_prompt_result_eval_scores")
    op.drop_index(
        op.f("ix_prompt_experiment_test_case_prompt_results_test_case_id"),
        table_name="prompt_experiment_test_case_prompt_results",
    )
    op.drop_index(
        op.f("ix_prompt_experiment_test_case_prompt_results_id"),
        table_name="prompt_experiment_test_case_prompt_results",
    )
    op.drop_table("prompt_experiment_test_case_prompt_results")
    op.drop_index(
        op.f("ix_prompt_experiment_test_cases_status"),
        table_name="prompt_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_prompt_experiment_test_cases_id"),
        table_name="prompt_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_prompt_experiment_test_cases_experiment_id"),
        table_name="prompt_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_prompt_experiment_test_cases_dataset_row_id"),
        table_name="prompt_experiment_test_cases",
    )
    op.drop_table("prompt_experiment_test_cases")
    op.drop_index(
        op.f("ix_prompt_experiments_task_id"), table_name="prompt_experiments"
    )
    op.drop_index(op.f("ix_prompt_experiments_status"), table_name="prompt_experiments")
    op.drop_index(
        op.f("ix_prompt_experiments_prompt_name"), table_name="prompt_experiments"
    )
    op.drop_index(op.f("ix_prompt_experiments_id"), table_name="prompt_experiments")
    op.drop_index(
        op.f("ix_prompt_experiments_created_at"), table_name="prompt_experiments"
    )
    op.drop_table("prompt_experiments")
    # ### end Alembic commands ###
