"""Add contractor_id to invoices table

Revision ID: add_contractor_id_to_invoices
Revises: complete_schema_migration
Create Date: 2025-08-22 12:13:00.000000

"""
from alembic import op
import sqlalchemy as sa

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

def upgrade():
    # Add contractor_id column to invoices table
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    
    if 'invoices' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('invoices')]
        
        if 'contractor_id' not in columns:
            op.add_column('invoices', sa.Column('contractor_id', sa.Integer(), nullable=False))
            op.create_foreign_key('fk_invoices_contractor_id', 'invoices', 'users', ['contractor_id'], ['id'])

def downgrade():
    # Remove contractor_id column from invoices table
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    
    if 'invoices' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('invoices')]
        
        if 'contractor_id' in columns:
            op.drop_constraint('fk_invoices_contractor_id', 'invoices', type_='foreignkey')
            op.drop_column('invoices', 'contractor_id')