from alembic import op
import sqlalchemy as sa
import sqlmodel.sql.sqltypes
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = '20f584dc80d2'
down_revision = '38a9c73bfce2'
branch_labels = None
depends_on = None

def upgrade():    
    # Create new tables
    op.create_table('checkpoint_blobs',
        sa.Column('thread_id', sa.Uuid(), nullable=False),
        sa.Column('checkpoint_ns', sa.String(), nullable=False, server_default=''),
        sa.Column('channel', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column('version', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column('type', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column('blob', sa.LargeBinary(), nullable=True),
        sa.ForeignKeyConstraint(['thread_id'], ['thread.id']),
        sa.PrimaryKeyConstraint('thread_id', 'checkpoint_ns', 'channel', 'version')
    )
    op.create_table('checkpoint_writes',
        sa.Column('thread_id', sa.Uuid(), nullable=False),
        sa.Column('checkpoint_ns', sa.String(), nullable=False, server_default=''),
        sa.Column('checkpoint_id', sa.Uuid(), nullable=False),
        sa.Column('task_id', sa.Uuid(), nullable=False),
        sa.Column('idx', sa.Integer(), nullable=False),
        sa.Column('channel', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column('type', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column('blob', sa.LargeBinary(), nullable=False),
        sa.ForeignKeyConstraint(['thread_id'], ['thread.id']),
        sa.PrimaryKeyConstraint('thread_id', 'checkpoint_ns', 'checkpoint_id', 'task_id', 'idx')
    )
    
    # Drop the old table
    op.drop_table('writes')

    # Rename and recreate the checkpoints table
    op.rename_table('checkpoints', 'checkpoints_old')
    
    op.create_table(
        'checkpoints',
        sa.Column('thread_id', sa.Uuid(), nullable=False),
        sa.Column('checkpoint_ns', sa.String(), nullable=False, server_default=''),
        sa.Column('checkpoint_id', sa.Uuid(), nullable=False),
        sa.Column('parent_checkpoint_id', sa.Uuid(), nullable=True),
        sa.Column('type', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column('checkpoint', postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'),
        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
        sa.ForeignKeyConstraint(['thread_id'], ['thread.id']),
        sa.PrimaryKeyConstraint('thread_id', 'checkpoint_ns', 'checkpoint_id')
    )
    
    # Drop the old checkpoints table
    op.drop_table('checkpoints_old')
    
    # Clear the threads table
    op.execute('DELETE FROM thread')
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    
    # Recreate the old checkpoints table
    op.create_table(
        'checkpoints_old',
        sa.Column('thread_id', sa.Uuid(), nullable=False),
        sa.Column('thread_ts', sa.Uuid(), nullable=False),
        sa.Column('parent_ts', sa.Uuid(), nullable=True),
        sa.Column('checkpoint', sa.LargeBinary(), nullable=True),
        sa.Column('metadata', sa.LargeBinary(), nullable=False, server_default=sa.text("'\\x'::bytea")),
        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()),
        sa.ForeignKeyConstraint(['thread_id'], ['thread.id']),
        sa.PrimaryKeyConstraint('thread_id', 'thread_ts')
    )
    
    # Drop the new checkpoints table
    op.drop_table('checkpoints')
    
    # Rename the old table back to 'checkpoints'
    op.rename_table('checkpoints_old', 'checkpoints')
    
    # Recreate the old 'writes' table
    op.create_table('writes',
        sa.Column('thread_id', sa.UUID(), autoincrement=False, nullable=False),
        sa.Column('thread_ts', sa.UUID(), autoincrement=False, nullable=False),
        sa.Column('task_id', sa.UUID(), autoincrement=False, nullable=False),
        sa.Column('idx', sa.INTEGER(), autoincrement=False, nullable=False),
        sa.Column('channel', sa.VARCHAR(), autoincrement=False, nullable=False),
        sa.Column('value', postgresql.BYTEA(), autoincrement=False, nullable=False),
        sa.ForeignKeyConstraint(['thread_id'], ['thread.id'], name='writes_thread_id_fkey'),
        sa.PrimaryKeyConstraint('thread_id', 'thread_ts', 'task_id', 'idx', name='writes_pkey')
    )
    
    # Drop the new tables
    op.drop_table('checkpoint_writes')
    op.drop_table('checkpoint_blobs')
    
    # Clear the threads table
    op.execute('DELETE FROM thread')
    # ### end Alembic commands ###