"""Add transfer tracking fields to Payment model

Revision ID: 4bf0f7199242
Revises: f9a8b2c3d4e5
Create Date: 2025-10-24 00:40:41.026315

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = '4bf0f7199242'
down_revision = 'f9a8b2c3d4e5'
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    existing_tables = inspector.get_table_names()
    
    # Drop user_sessions table if it exists
    if 'user_sessions' in existing_tables:
        op.drop_table('user_sessions')
    
    # Add columns to payments table if they don't exist
    if 'payments' in existing_tables:
        columns = [col['name'] for col in inspector.get_columns('payments')]
        with op.batch_alter_table('payments', schema=None) as batch_op:
            if 'stripe_transfer_id' not in columns:
                batch_op.add_column(sa.Column('stripe_transfer_id', sa.String(length=255), nullable=True))
            if 'stripe_transfer_status' not in columns:
                batch_op.add_column(sa.Column('stripe_transfer_status', sa.String(length=50), nullable=True))
            if 'stripe_transfer_id' not in columns:  # Only create constraint if column was added
                batch_op.create_unique_constraint(None, ['stripe_transfer_id'])

    # Drop time_entries indexes if they exist
    if 'time_entries' in existing_tables:
        existing_indexes = {idx['name'] for idx in inspector.get_indexes('time_entries')}
        with op.batch_alter_table('time_entries', schema=None) as batch_op:
            if 'idx_time_entries_contract_approval' in existing_indexes:
                batch_op.drop_index(batch_op.f('idx_time_entries_contract_approval'))
            if 'idx_time_entries_contract_pending' in existing_indexes:
                batch_op.drop_index(batch_op.f('idx_time_entries_contract_pending'))
            if 'idx_time_entries_date_range' in existing_indexes:
                batch_op.drop_index(batch_op.f('idx_time_entries_date_range'))
            if 'idx_time_entries_worker_active' in existing_indexes:
                batch_op.drop_index(batch_op.f('idx_time_entries_worker_active'))
            if 'idx_time_entries_worker_date' in existing_indexes:
                batch_op.drop_index(batch_op.f('idx_time_entries_worker_date'))

    # Drop users columns if they exist
    if 'users' in existing_tables:
        columns = [col['name'] for col in inspector.get_columns('users')]
        columns_to_drop = ['hourly_rate', 'stripe_customer_id', 'certifications', 'reset_token_expiry',
                          'address', 'abn', 'is_verified', 'rating', 'verification_token_expiry',
                          'portfolio_url', 'two_factor_enabled', 'bio', 'reset_token', 'two_factor_secret',
                          'company_name', 'license_number', 'skills', 'insurance_details', 'profile_picture',
                          'years_experience', 'completed_jobs', 'availability_status', 'verification_token', 'phone']
        
        with op.batch_alter_table('users', schema=None) as batch_op:
            for col in columns_to_drop:
                if col in columns:
                    batch_op.drop_column(col)

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table('users', schema=None) as batch_op:
        batch_op.add_column(sa.Column('phone', sa.VARCHAR(length=20), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('verification_token', sa.VARCHAR(length=100), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('availability_status', sa.VARCHAR(length=50), server_default=sa.text("'available'::character varying"), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('completed_jobs', sa.INTEGER(), server_default=sa.text('0'), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('years_experience', sa.INTEGER(), server_default=sa.text('0'), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('profile_picture', sa.VARCHAR(length=500), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('insurance_details', sa.TEXT(), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('skills', sa.TEXT(), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('license_number', sa.VARCHAR(length=100), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('company_name', sa.VARCHAR(length=200), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('two_factor_secret', sa.VARCHAR(length=32), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('reset_token', sa.VARCHAR(length=100), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('bio', sa.TEXT(), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('two_factor_enabled', sa.BOOLEAN(), server_default=sa.text('false'), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('portfolio_url', sa.VARCHAR(length=500), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('verification_token_expiry', postgresql.TIMESTAMP(), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('rating', sa.NUMERIC(precision=3, scale=2), server_default=sa.text('0.00'), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('is_verified', sa.BOOLEAN(), server_default=sa.text('false'), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('abn', sa.VARCHAR(length=20), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('address', sa.TEXT(), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('reset_token_expiry', postgresql.TIMESTAMP(), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('certifications', sa.TEXT(), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('stripe_customer_id', sa.VARCHAR(length=255), autoincrement=False, nullable=True))
        batch_op.add_column(sa.Column('hourly_rate', sa.NUMERIC(precision=10, scale=2), autoincrement=False, nullable=True))

    with op.batch_alter_table('time_entries', schema=None) as batch_op:
        batch_op.create_index(batch_op.f('idx_time_entries_worker_date'), ['worker_id', 'work_date'], unique=False)
        batch_op.create_index(batch_op.f('idx_time_entries_worker_active'), ['worker_id', 'work_date', 'clock_out'], unique=False)
        batch_op.create_index(batch_op.f('idx_time_entries_date_range'), ['work_date', 'contract_id'], unique=False)
        batch_op.create_index(batch_op.f('idx_time_entries_contract_pending'), ['contract_id', 'is_approved', 'work_date'], unique=False)
        batch_op.create_index(batch_op.f('idx_time_entries_contract_approval'), ['contract_id', 'is_approved'], unique=False)

    with op.batch_alter_table('payments', schema=None) as batch_op:
        batch_op.drop_constraint(None, type_='unique')
        batch_op.drop_column('stripe_transfer_status')
        batch_op.drop_column('stripe_transfer_id')

    op.create_table('user_sessions',
    sa.Column('id', sa.INTEGER(), autoincrement=True, nullable=False),
    sa.Column('user_id', sa.INTEGER(), autoincrement=False, nullable=False),
    sa.Column('session_token', sa.VARCHAR(length=255), autoincrement=False, nullable=False),
    sa.Column('ip_address', sa.VARCHAR(length=45), autoincrement=False, nullable=True),
    sa.Column('user_agent', sa.TEXT(), autoincrement=False, nullable=True),
    sa.Column('created_at', postgresql.TIMESTAMP(), server_default=sa.text('now()'), autoincrement=False, nullable=True),
    sa.Column('expires_at', postgresql.TIMESTAMP(), autoincrement=False, nullable=False),
    sa.Column('is_active', sa.BOOLEAN(), server_default=sa.text('true'), autoincrement=False, nullable=True),
    sa.ForeignKeyConstraint(['user_id'], ['users.id'], name=op.f('user_sessions_user_id_fkey'), ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('id', name=op.f('user_sessions_pkey')),
    sa.UniqueConstraint('session_token', name=op.f('user_sessions_session_token_key'))
    )
    # ### end Alembic commands ###