"""Add rule result token columns

Revision ID: ba03ca11c6ab
Revises: 769fc709b3a7
Create Date: 2023-09-07 12:53:55.095826

"""

import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "ba03ca11c6ab"
down_revision = "769fc709b3a7"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_unique_constraint(None, "configurations", ["name"])
    op.create_unique_constraint(
        None,
        "inference_prompt_contents",
        ["inference_prompt_id"],
    )
    op.create_unique_constraint(
        None,
        "inference_response_contents",
        ["inference_response_id"],
    )
    op.add_column(
        "prompt_rule_results",
        sa.Column("user_input_tokens", sa.Integer(), nullable=True),
    )
    op.add_column(
        "prompt_rule_results",
        sa.Column("prompt_tokens", sa.Integer(), nullable=True),
    )
    op.add_column(
        "prompt_rule_results",
        sa.Column("completion_tokens", sa.Integer(), nullable=True),
    )
    op.add_column(
        "response_rule_results",
        sa.Column("user_input_tokens", sa.Integer(), nullable=True),
    )
    op.add_column(
        "response_rule_results",
        sa.Column("prompt_tokens", sa.Integer(), nullable=True),
    )
    op.add_column(
        "response_rule_results",
        sa.Column("completion_tokens", sa.Integer(), nullable=True),
    )

    op.execute(
        "UPDATE prompt_rule_results SET user_input_tokens=0 WHERE user_input_tokens IS null",
    )
    op.execute(
        "UPDATE prompt_rule_results SET prompt_tokens=0 WHERE prompt_tokens IS null",
    )
    op.execute(
        "UPDATE prompt_rule_results SET completion_tokens=0 WHERE completion_tokens IS null",
    )

    op.execute(
        "UPDATE response_rule_results SET user_input_tokens=0 WHERE user_input_tokens IS null",
    )
    op.execute(
        "UPDATE response_rule_results SET prompt_tokens=0 WHERE prompt_tokens IS null",
    )
    op.execute(
        "UPDATE response_rule_results SET completion_tokens=0 WHERE completion_tokens IS null",
    )

    op.alter_column(
        "prompt_rule_results",
        sa.Column("user_input_tokens", sa.Integer(), nullable=False),
    )
    op.alter_column(
        "prompt_rule_results",
        sa.Column("prompt_tokens", sa.Integer(), nullable=False),
    )
    op.alter_column(
        "prompt_rule_results",
        sa.Column("completion_tokens", sa.Integer(), nullable=False),
    )
    op.alter_column(
        "response_rule_results",
        sa.Column("user_input_tokens", sa.Integer(), nullable=False),
    )
    op.alter_column(
        "response_rule_results",
        sa.Column("prompt_tokens", sa.Integer(), nullable=False),
    )
    op.alter_column(
        "response_rule_results",
        sa.Column("completion_tokens", sa.Integer(), nullable=False),
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("response_rule_results", "completion_tokens")
    op.drop_column("response_rule_results", "prompt_tokens")
    op.drop_column("response_rule_results", "user_input_tokens")
    op.drop_column("prompt_rule_results", "completion_tokens")
    op.drop_column("prompt_rule_results", "prompt_tokens")
    op.drop_column("prompt_rule_results", "user_input_tokens")
    # ### end Alembic commands ###
