"""

Revision ID: 111111111111
Revises: 000000000000
Create Date: 2018-09-25 14:54:16.687076

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '111111111111'
down_revision = '000000000000'
branch_labels = None
depends_on = None


# Which tables have a column with this type that needs to be migrated
AFFECTED = [
    ('table1', 'column1'),
    ('table1', 'column2'),
    ('table2', 'column'),
]

# Enum choices, they need to be hard-coded here
TYPE_NAME = 'color_enum'
OLD_OPTIONS = ['red', 'blue', 'yellow']
NEW_OPTIONS = ['red', 'blue', 'green']

# Conversion from old values to new ones
CONVERSION_FORWARD = {
    'yellow': 'green',
    'blue': 'green',
}
# Conversion from new values to old ones
CONVERSION_BACKWARD = {
    'green': 'yellow',
}


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # ### end Alembic commands ###
    new_type = sa.Enum(*NEW_OPTIONS, name=TYPE_NAME + '_new')
    old_type = sa.Enum(*OLD_OPTIONS, name=TYPE_NAME)

    # Create the new ENUM
    new_type.create(op.get_bind(), checkfirst=False)

    # Convert the Data
    for table_name, column_name in AFFECTED:
        # Create the new column to copy the data to
        op.add_column(table_name, sa.Column(column_name + '_new', new_type, nullable=True))
        # Copy the data over, doing any necessary conversions
        op.execute("""
            UPDATE {table_name} SET {column_name}_new=CASE(
                {conversions}
                ELSE {column_name}::text
            )::{new_type_name}
        """.format(
            table_name=table_name,
            new_type_name=TYPE_NAME + '_new',
            conversions="\n".join(
                "WHEN {column_name}='{old_val}' THEN '{new_val}'".format(
                    old_val=old_val,
                    new_val=new_val,
                    column_name=column_name,
                ) for old_val, new_val in CONVERSION_FORWARD.items()
            ),
        ))
        # Remove the old Column
        op.drop_column(table_name, column_name)
        # And rename the new one into its place
        op.alter_column(table_name, column_name + '_new', new_column_name=column_name)

    # Remove the old type
    old_type.drop(op.get_bind(), checkfirst=False)
    op.execute("ALTER TYPE {type_name}_new RENAME TO {type_name}".format(
        type_name=TYPE_NAME,
    ))


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # ### end Alembic commands ###
    new_type = sa.Enum(*NEW_OPTIONS, name=TYPE_NAME)
    old_type = sa.Enum(*OLD_OPTIONS, name=TYPE_NAME + '_old')

    # re-Create the old ENUM
    old_type.create(op.get_bind(), checkfirst=False)

    # Convert the Data
    for table_name, column_name in AFFECTED:
        op.add_column(table_name, sa.Column(column_name + '_old', old_type, nullable=True))

        op.execute("""
            UPDATE {table_name} SET {column_name}_old=CASE(
                {conversions}
                ELSE {column_name}::text
            )::{old_type_name}
        """.format(
            table_name=table_name,
            old_type_name=TYPE_NAME + '_old',
            conversions="\n".join(
                "WHEN {column_name}='{old_val}' THEN '{new_val}'".format(
                    old_val=old_val,
                    new_val=new_val,
                    column_name=column_name,
                ) for old_val, new_val in CONVERSION_BACKWARD.items()
            ),
        ))
        op.drop_column(table_name, column_name)
        op.alter_column(table_name, column_name + '_old', new_column_name=column_name)

    # Remove the old type
    new_type.drop(op.get_bind(), checkfirst=False)
    op.execute("ALTER TYPE {type_name}_old RENAME TO {type_name}".format(
        type_name=TYPE_NAME,
    ))
