"""create rag experiments

Revision ID: 522b83a56135
Revises: 44cd9fd55df9
Create Date: 2025-12-10 11:20:04.056258

"""

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

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


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "rag_experiments",
        sa.Column(
            "rag_configs",
            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("notebook_id", sa.String(), nullable=True),
        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(["notebook_id"], ["notebooks.id"], ondelete="SET NULL"),
        sa.ForeignKeyConstraint(["task_id"], ["tasks.id"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_rag_experiments_created_at"),
        "rag_experiments",
        ["created_at"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiments_id"),
        "rag_experiments",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiments_notebook_id"),
        "rag_experiments",
        ["notebook_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiments_status"),
        "rag_experiments",
        ["status"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiments_task_id"),
        "rag_experiments",
        ["task_id"],
        unique=False,
    )
    op.create_table(
        "rag_experiment_test_cases",
        sa.Column("experiment_id", sa.String(), 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"],
            ["rag_experiments.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_rag_experiment_test_cases_dataset_row_id"),
        "rag_experiment_test_cases",
        ["dataset_row_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiment_test_cases_experiment_id"),
        "rag_experiment_test_cases",
        ["experiment_id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiment_test_cases_id"),
        "rag_experiment_test_cases",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiment_test_cases_status"),
        "rag_experiment_test_cases",
        ["status"],
        unique=False,
    )
    op.create_table(
        "rag_experiment_test_case_rag_results",
        sa.Column("id", sa.String(), nullable=False),
        sa.Column("test_case_id", sa.String(), nullable=False),
        sa.Column("rag_config_key", sa.String(), nullable=False),
        sa.Column("rag_config_type", sa.String(), nullable=False),
        sa.Column("setting_configuration_id", sa.UUID(), nullable=True),
        sa.Column("setting_configuration_version", sa.Integer(), nullable=True),
        sa.Column("query_text", sa.Text(), nullable=False),
        sa.Column(
            "search_output",
            postgresql.JSON(astext_type=sa.Text()),
            nullable=False,
        ),
        sa.Column("created_at", sa.TIMESTAMP(), nullable=False),
        sa.ForeignKeyConstraint(
            ["test_case_id"],
            ["rag_experiment_test_cases.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_rag_experiment_test_case_rag_results_id"),
        "rag_experiment_test_case_rag_results",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiment_test_case_rag_results_rag_config_key"),
        "rag_experiment_test_case_rag_results",
        ["rag_config_key"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiment_test_case_rag_results_test_case_id"),
        "rag_experiment_test_case_rag_results",
        ["test_case_id"],
        unique=False,
    )
    op.create_table(
        "rag_experiment_test_case_rag_result_eval_scores",
        sa.Column("rag_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(
            ["rag_result_id"],
            ["rag_experiment_test_case_rag_results.id"],
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_rag_experiment_test_case_rag_result_eval_scores_id"),
        "rag_experiment_test_case_rag_result_eval_scores",
        ["id"],
        unique=False,
    )
    op.create_index(
        op.f("ix_rag_experiment_test_case_rag_result_eval_scores_rag_result_id"),
        "rag_experiment_test_case_rag_result_eval_scores",
        ["rag_result_id"],
        unique=False,
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(
        op.f("ix_rag_experiment_test_case_rag_result_eval_scores_rag_result_id"),
        table_name="rag_experiment_test_case_rag_result_eval_scores",
    )
    op.drop_index(
        op.f("ix_rag_experiment_test_case_rag_result_eval_scores_id"),
        table_name="rag_experiment_test_case_rag_result_eval_scores",
    )
    op.drop_table("rag_experiment_test_case_rag_result_eval_scores")
    op.drop_index(
        op.f("ix_rag_experiment_test_case_rag_results_test_case_id"),
        table_name="rag_experiment_test_case_rag_results",
    )
    op.drop_index(
        op.f("ix_rag_experiment_test_case_rag_results_rag_config_key"),
        table_name="rag_experiment_test_case_rag_results",
    )
    op.drop_index(
        op.f("ix_rag_experiment_test_case_rag_results_id"),
        table_name="rag_experiment_test_case_rag_results",
    )
    op.drop_table("rag_experiment_test_case_rag_results")
    op.drop_index(
        op.f("ix_rag_experiment_test_cases_status"),
        table_name="rag_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_rag_experiment_test_cases_id"),
        table_name="rag_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_rag_experiment_test_cases_experiment_id"),
        table_name="rag_experiment_test_cases",
    )
    op.drop_index(
        op.f("ix_rag_experiment_test_cases_dataset_row_id"),
        table_name="rag_experiment_test_cases",
    )
    op.drop_table("rag_experiment_test_cases")
    op.drop_index(op.f("ix_rag_experiments_task_id"), table_name="rag_experiments")
    op.drop_index(op.f("ix_rag_experiments_status"), table_name="rag_experiments")
    op.drop_index(op.f("ix_rag_experiments_notebook_id"), table_name="rag_experiments")
    op.drop_index(op.f("ix_rag_experiments_id"), table_name="rag_experiments")
    op.drop_index(op.f("ix_rag_experiments_created_at"), table_name="rag_experiments")
    op.drop_table("rag_experiments")
    # ### end Alembic commands ###
