"""agentic experiment models

Revision ID: 06b0610b582c
Revises: 8aba08c2934f
Create Date: 2026-01-05 12:27:37.340393

"""

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

# revision identifiers, used by Alembic.
revision = "06b0610b582c"
down_revision = "8aba08c2934f"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "agentic_experiments",
        sa.Column(
            "http_template",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column(
            "template_variable_mapping",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        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("dataset_id", sa.UUID(), nullable=False),
        sa.Column("dataset_version", sa.Integer(), nullable=False),
        sa.Column(
            "dataset_row_filter",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=True,
        ),
        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(
            ["dataset_id", "dataset_version"],
            ["dataset_versions.dataset_id", "dataset_versions.version_number"],
        ),
        sa.ForeignKeyConstraint(["task_id"], ["tasks.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_agentic_experiments_created_at"),
        "agentic_experiments",
        ["created_at"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentic_experiments_id"),
        "agentic_experiments",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentic_experiments_status"),
        "agentic_experiments",
        ["status"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentic_experiments_task_id"),
        "agentic_experiments",
        ["task_id"],
        unique=False,
    )
    op.create_table(
        "agentic_experiment_test_cases",
        sa.Column("experiment_id", sa.String(), nullable=False),
        sa.Column(
            "template_input_variables",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("status", sa.String(), nullable=False),
        sa.Column("dataset_row_id", sa.String(), 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"],
            ["agentic_experiments.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_agentic_experiment_test_cases_dataset_row_id"),
        "agentic_experiment_test_cases",
        ["dataset_row_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentic_experiment_test_cases_experiment_id"),
        "agentic_experiment_test_cases",
        ["experiment_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentic_experiment_test_cases_id"),
        "agentic_experiment_test_cases",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentic_experiment_test_cases_status"),
        "agentic_experiment_test_cases",
        ["status"],
        unique=False,
    )
    op.create_table(
        "agentic_experiment_test_case_agentic_results",
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("test_case_id", sa.String(), nullable=False),
        sa.Column("request_url", sa.Text(), nullable=False),
        sa.Column(
            "request_headers",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column(
            "request_body",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column(
            "response_output",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.ForeignKeyConstraint(
            ["test_case_id"],
            ["agentic_experiment_test_cases.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_agentic_experiment_test_case_agentic_results_id"),
        "agentic_experiment_test_case_agentic_results",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentic_experiment_test_case_agentic_results_test_case_id"),
        "agentic_experiment_test_case_agentic_results",
        ["test_case_id"],
        unique=True,
    )
    op.create_table(
        "agentic_experiment_test_case_agentic_result_eval_scores",
        sa.Column("agentic_result_id", sa.String(), nullable=False),
        sa.Column("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(
            ["agentic_result_id"],
            ["agentic_experiment_test_case_agentic_results.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f(
            "ix_agentic_experiment_test_case_agentic_result_eval_scores_agentic_result_id",
        ),
        "agentic_experiment_test_case_agentic_result_eval_scores",
        ["agentic_result_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_agentic_experiment_test_case_agentic_result_eval_scores_id"),
        "agentic_experiment_test_case_agentic_result_eval_scores",
        ["id"],
        unique=False,
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(
        op.f("ix_agentic_experiment_test_case_agentic_result_eval_scores_id"),
        table_name="agentic_experiment_test_case_agentic_result_eval_scores",
    )
    op.drop_index(
        op.f(
            "ix_agentic_experiment_test_case_agentic_result_eval_scores_agentic_result_id",
        ),
        table_name="agentic_experiment_test_case_agentic_result_eval_scores",
    )
    op.drop_table("agentic_experiment_test_case_agentic_result_eval_scores")
    op.drop_index(
        op.f("ix_agentic_experiment_test_case_agentic_results_test_case_id"),
        table_name="agentic_experiment_test_case_agentic_results",
    )
    op.drop_index(
        op.f("ix_agentic_experiment_test_case_agentic_results_id"),
        table_name="agentic_experiment_test_case_agentic_results",
    )
    op.drop_table("agentic_experiment_test_case_agentic_results")
    op.drop_index(
        op.f("ix_agentic_experiment_test_cases_status"),
        table_name="agentic_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_agentic_experiment_test_cases_id"),
        table_name="agentic_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_agentic_experiment_test_cases_experiment_id"),
        table_name="agentic_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_agentic_experiment_test_cases_dataset_row_id"),
        table_name="agentic_experiment_test_cases",
    )
    op.drop_table("agentic_experiment_test_cases")
    op.drop_index(
        op.f("ix_agentic_experiments_task_id"),
        table_name="agentic_experiments",
    )
    op.drop_index(
        op.f("ix_agentic_experiments_status"),
        table_name="agentic_experiments",
    )
    op.drop_index(op.f("ix_agentic_experiments_id"), table_name="agentic_experiments")
    op.drop_index(
        op.f("ix_agentic_experiments_created_at"),
        table_name="agentic_experiments",
    )
    op.drop_table("agentic_experiments")
    # ### end Alembic commands ###
