"""Add the table "trace_call"

Revision ID: f3edea6dbde2
Revises: 2cf3fc3d3dc5
Create Date: 2024-06-10 09:42:36.477087

"""
import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = 'f3edea6dbde2'
down_revision = '2cf3fc3d3dc5'
branch_labels = None
depends_on = None


def upgrade():
    # Disable the foreign key constraints during the migration.
    # https://alembic.sqlalchemy.org/en/latest/batch.html#dealing-with-referencing-foreign-keys
    op.execute('PRAGMA foreign_keys=OFF;')

    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        'trace_call',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('trace_call_no', sa.Integer(), nullable=False),
        sa.Column('started_at', sa.DateTime(), nullable=False),
        sa.Column('file_name', sa.String(), nullable=True),
        sa.Column('line_no', sa.Integer(), nullable=True),
        sa.Column('event', sa.String(), nullable=False),
        sa.Column('ended_at', sa.DateTime(), nullable=True),
        sa.Column('run_id', sa.Integer(), nullable=False),
        sa.Column('trace_id', sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(
            ['run_id'], ['run.id'], name=op.f('fk_trace_call_run_id_run')
        ),
        sa.ForeignKeyConstraint(
            ['trace_id'], ['trace.id'], name=op.f('fk_trace_call_trace_id_trace')
        ),
        sa.PrimaryKeyConstraint('id', name=op.f('pk_trace_call')),
        sa.UniqueConstraint(
            'run_id', 'trace_call_no', name=op.f('uq_trace_call_run_id')
        ),
    )
    with op.batch_alter_table('trace_call', schema=None) as batch_op:
        batch_op.create_index(batch_op.f('ix_trace_call_id'), ['id'], unique=False)

    with op.batch_alter_table('prompt', schema=None) as batch_op:
        batch_op.add_column(sa.Column('trace_call_id', sa.Integer(), nullable=True))
        batch_op.create_foreign_key(
            batch_op.f('fk_prompt_trace_call_id_trace_call'),
            'trace_call',
            ['trace_call_id'],
            ['id'],
        )

    # ### end Alembic commands ###

    # Re-enable the foreign key constraints
    op.execute('PRAGMA foreign_keys=ON;')


def downgrade():
    op.execute('PRAGMA foreign_keys=OFF;')

    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('prompt', schema=None) as batch_op:
        batch_op.drop_constraint(
            batch_op.f('fk_prompt_trace_call_id_trace_call'), type_='foreignkey'
        )
        batch_op.drop_column('trace_call_id')

    with op.batch_alter_table('trace_call', schema=None) as batch_op:
        batch_op.drop_index(batch_op.f('ix_trace_call_id'))

    op.drop_table('trace_call')
    # ### end Alembic commands ###

    op.execute('PRAGMA foreign_keys=ON;')
