"""Add dispute resolution agreement fields

Revision ID: e1588d3f77c7
Revises: add_accepted_workers_count
Create Date: 2025-10-16 00:52:40.839950

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = 'e1588d3f77c7'
down_revision = 'add_accepted_workers_count'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    
    # Check if contracts table exists and get columns
    if 'contracts' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('contracts')]
        
        # Only add columns if they don't exist
        columns_to_add = []
        if 'worker_agreed_resolve' not in columns:
            columns_to_add.append('worker_agreed_resolve')
        if 'contractor_agreed_resolve' not in columns:
            columns_to_add.append('contractor_agreed_resolve')
        
        if columns_to_add:
            # First add the columns as nullable
            with op.batch_alter_table('contracts', schema=None) as batch_op:
                if 'worker_agreed_resolve' in columns_to_add:
                    batch_op.add_column(sa.Column('worker_agreed_resolve', sa.Boolean(), nullable=True))
                if 'contractor_agreed_resolve' in columns_to_add:
                    batch_op.add_column(sa.Column('contractor_agreed_resolve', sa.Boolean(), nullable=True))
            
            # Update existing rows to have the default value of False
            connection = op.get_bind()
            update_parts = []
            if 'worker_agreed_resolve' in columns_to_add:
                update_parts.append('worker_agreed_resolve = false')
            if 'contractor_agreed_resolve' in columns_to_add:
                update_parts.append('contractor_agreed_resolve = false')
            
            if update_parts:
                connection.execute(sa.text(
                    f'UPDATE contracts SET {", ".join(update_parts)}'
                ))

            # Now alter the columns to be non-nullable
            with op.batch_alter_table('contracts', schema=None) as batch_op:
                if 'worker_agreed_resolve' in columns_to_add:
                    batch_op.alter_column('worker_agreed_resolve',
                           existing_type=sa.Boolean(),
                           nullable=False)
                if 'contractor_agreed_resolve' in columns_to_add:
                    batch_op.alter_column('contractor_agreed_resolve',
                           existing_type=sa.Boolean(),
                           nullable=False)

    # Handle jobs table column alteration
    if 'jobs' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('jobs')]
        if 'accepted_workers_count' in columns:
            with op.batch_alter_table('jobs', schema=None) as batch_op:
                batch_op.alter_column('accepted_workers_count',
                       existing_type=sa.INTEGER(),
                       nullable=True,
                       existing_server_default=sa.text('0'))

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    
    if 'jobs' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('jobs')]
        if 'accepted_workers_count' in columns:
            with op.batch_alter_table('jobs', schema=None) as batch_op:
                batch_op.alter_column('accepted_workers_count',
                       existing_type=sa.INTEGER(),
                       nullable=False,
                       existing_server_default=sa.text('0'))

    if 'contracts' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('contracts')]
        
        with op.batch_alter_table('contracts', schema=None) as batch_op:
            if 'contractor_agreed_resolve' in columns:
                batch_op.drop_column('contractor_agreed_resolve')
            if 'worker_agreed_resolve' in columns:
                batch_op.drop_column('worker_agreed_resolve')

    # ### end Alembic commands ###