"""Add Stripe Connect fields to users table

Revision ID: add_stripe_connect_to_users
Revises: add_contractor_id_to_invoices
Create Date: 2025-08-22 12:14:00.000000

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = 'add_stripe_connect_to_users'
down_revision = 'add_contractor_id_to_invoices'
branch_labels = None
depends_on = None

def upgrade():
    # Add Stripe Connect columns to users table
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    
    if 'users' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('users')]
        
        columns_to_add = []
        if 'stripe_account_id' not in columns:
            columns_to_add.append(('stripe_account_id', sa.Column('stripe_account_id', sa.String(255), nullable=True)))
        if 'stripe_onboarding_complete' not in columns:
            columns_to_add.append(('stripe_onboarding_complete', sa.Column('stripe_onboarding_complete', sa.Boolean(), nullable=False, server_default='false')))
        if 'stripe_payouts_enabled' not in columns:
            columns_to_add.append(('stripe_payouts_enabled', sa.Column('stripe_payouts_enabled', sa.Boolean(), nullable=False, server_default='false')))
        if 'stripe_charges_enabled' not in columns:
            columns_to_add.append(('stripe_charges_enabled', sa.Column('stripe_charges_enabled', sa.Boolean(), nullable=False, server_default='false')))
        if 'stripe_account_created_at' not in columns:
            columns_to_add.append(('stripe_account_created_at', sa.Column('stripe_account_created_at', sa.DateTime(), nullable=True)))
        
        for col_name, col_def in columns_to_add:
            op.add_column('users', col_def)
        
        # Create unique constraint for stripe_account_id if it was added
        if 'stripe_account_id' in [col[0] for col in columns_to_add]:
            op.create_unique_constraint('uq_users_stripe_account_id', 'users', ['stripe_account_id'])

def downgrade():
    # Remove Stripe Connect columns from users table
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    
    if 'users' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('users')]
        
        try:
            op.drop_constraint('uq_users_stripe_account_id', 'users', type_='unique')
        except:
            pass
        
        if 'stripe_account_created_at' in columns:
            op.drop_column('users', 'stripe_account_created_at')
        if 'stripe_charges_enabled' in columns:
            op.drop_column('users', 'stripe_charges_enabled')
        if 'stripe_payouts_enabled' in columns:
            op.drop_column('users', 'stripe_payouts_enabled')
        if 'stripe_onboarding_complete' in columns:
            op.drop_column('users', 'stripe_onboarding_complete')
        if 'stripe_account_id' in columns:
            op.drop_column('users', 'stripe_account_id')