"""Add company_id to batch_uploads table.

Revision ID: 20260109_batch_company
Revises: 20251231_add_original_filename_to_metrics
Create Date: 2026-01-09

This migration adds the company_id column to batch_uploads to support
automatic receipt creation during bulk upload processing.
"""

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '20260109_batch_company'
down_revision = None  # Will be auto-detected
branch_labels = None
depends_on = None


def upgrade() -> None:
    """Add company_id column to batch_uploads table."""
    # Check if column already exists (SQLModel may have created it)
    conn = op.get_bind()
    inspector = sa.inspect(conn)

    # Check if batch_uploads table exists
    if 'batch_uploads' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('batch_uploads')]
        if 'company_id' not in columns:
            op.add_column(
                'batch_uploads',
                sa.Column('company_id', sa.Integer(), nullable=True)
            )
            # Create index for company_id
            op.create_index(
                'ix_batch_uploads_company_id',
                'batch_uploads',
                ['company_id'],
                unique=False
            )


def downgrade() -> None:
    """Remove company_id column from batch_uploads table."""
    conn = op.get_bind()
    inspector = sa.inspect(conn)

    if 'batch_uploads' in inspector.get_table_names():
        columns = [col['name'] for col in inspector.get_columns('batch_uploads')]
        if 'company_id' in columns:
            op.drop_index('ix_batch_uploads_company_id', table_name='batch_uploads')
            op.drop_column('batch_uploads', 'company_id')
