"""Add production stage enum and columns

Revision ID: 003
Revises: 002
Create Date: 2025-01-17 15:30:00.000000

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

# revision identifiers, used by Alembic.
revision = '003'
down_revision = '002'
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Create ProductionStage enum type
    op.execute("""
        CREATE TYPE productionstage AS ENUM (
            'raw_ingredient', 'par_baked', 'fully_baked',
            'prepared_dough', 'frozen_product'
        );
    """)

    # Add production_stage column to stock table
    op.add_column('stock', sa.Column('production_stage',
        sa.Enum('raw_ingredient', 'par_baked', 'fully_baked', 'prepared_dough', 'frozen_product', name='productionstage'),
        nullable=False, server_default='raw_ingredient'))

    # Add transformation_reference column to stock table
    op.add_column('stock', sa.Column('transformation_reference', sa.String(100), nullable=True))

    # Add stage-specific expiration tracking columns
    op.add_column('stock', sa.Column('original_expiration_date', sa.DateTime(timezone=True), nullable=True))
    op.add_column('stock', sa.Column('transformation_date', sa.DateTime(timezone=True), nullable=True))
    op.add_column('stock', sa.Column('final_expiration_date', sa.DateTime(timezone=True), nullable=True))

    # Create product_transformations table
    op.create_table(
        'product_transformations',
        sa.Column('id', postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column('tenant_id', postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column('transformation_reference', sa.String(100), nullable=False),
        sa.Column('source_ingredient_id', postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column('target_ingredient_id', postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column('source_stage', sa.Enum('raw_ingredient', 'par_baked', 'fully_baked', 'prepared_dough', 'frozen_product', name='productionstage'), nullable=False),
        sa.Column('target_stage', sa.Enum('raw_ingredient', 'par_baked', 'fully_baked', 'prepared_dough', 'frozen_product', name='productionstage'), nullable=False),
        sa.Column('source_quantity', sa.Float(), nullable=False),
        sa.Column('target_quantity', sa.Float(), nullable=False),
        sa.Column('conversion_ratio', sa.Float(), nullable=False, server_default='1.0'),
        sa.Column('expiration_calculation_method', sa.String(50), nullable=False, server_default='days_from_transformation'),
        sa.Column('expiration_days_offset', sa.Integer(), nullable=True),
        sa.Column('transformation_date', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
        sa.Column('process_notes', sa.Text(), nullable=True),
        sa.Column('performed_by', postgresql.UUID(as_uuid=True), nullable=True),
        sa.Column('source_batch_numbers', sa.Text(), nullable=True),
        sa.Column('target_batch_number', sa.String(100), nullable=True),
        sa.Column('is_completed', sa.Boolean(), nullable=True, server_default='true'),
        sa.Column('is_reversed', sa.Boolean(), nullable=True, server_default='false'),
        sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False),
        sa.Column('created_by', postgresql.UUID(as_uuid=True), nullable=True),
        sa.ForeignKeyConstraint(['source_ingredient_id'], ['ingredients.id'], ),
        sa.ForeignKeyConstraint(['target_ingredient_id'], ['ingredients.id'], ),
        sa.PrimaryKeyConstraint('id')
    )

    # Add new indexes for enhanced functionality
    op.create_index('idx_stock_production_stage', 'stock', ['tenant_id', 'production_stage', 'is_available'])
    op.create_index('idx_stock_transformation', 'stock', ['tenant_id', 'transformation_reference'])
    op.create_index('idx_stock_final_expiration', 'stock', ['tenant_id', 'final_expiration_date', 'is_available'])

    # Create indexes for product_transformations table
    op.create_index('idx_transformations_tenant_date', 'product_transformations', ['tenant_id', 'transformation_date'])
    op.create_index('idx_transformations_reference', 'product_transformations', ['transformation_reference'])
    op.create_index('idx_transformations_source', 'product_transformations', ['tenant_id', 'source_ingredient_id'])
    op.create_index('idx_transformations_target', 'product_transformations', ['tenant_id', 'target_ingredient_id'])
    op.create_index('idx_transformations_stages', 'product_transformations', ['source_stage', 'target_stage'])

    # Update existing stockmovementtype enum to include TRANSFORMATION
    op.execute("ALTER TYPE stockmovementtype ADD VALUE 'transformation';")


def downgrade() -> None:
    # Drop indexes for product_transformations
    op.drop_index('idx_transformations_stages', table_name='product_transformations')
    op.drop_index('idx_transformations_target', table_name='product_transformations')
    op.drop_index('idx_transformations_source', table_name='product_transformations')
    op.drop_index('idx_transformations_reference', table_name='product_transformations')
    op.drop_index('idx_transformations_tenant_date', table_name='product_transformations')

    # Drop new stock indexes
    op.drop_index('idx_stock_final_expiration', table_name='stock')
    op.drop_index('idx_stock_transformation', table_name='stock')
    op.drop_index('idx_stock_production_stage', table_name='stock')

    # Drop product_transformations table
    op.drop_table('product_transformations')

    # Remove new columns from stock table
    op.drop_column('stock', 'final_expiration_date')
    op.drop_column('stock', 'transformation_date')
    op.drop_column('stock', 'original_expiration_date')
    op.drop_column('stock', 'transformation_reference')
    op.drop_column('stock', 'production_stage')

    # Drop ProductionStage enum type
    op.execute("DROP TYPE productionstage;")

    # Note: Cannot easily remove 'transformation' from existing enum in PostgreSQL
    # This would require recreating the enum and updating all references
    # For now, we leave the enum value as it won't cause issues