'''add_rep_owner_keys_to_label

Revision ID: 7289d16beca6
Revises: 260e1763d6a0
Create Date: 2021-04-07 14:46:42.969405

'''
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '7289d16beca6'
down_revision = '1487dcffbdb4'
branch_labels = None
depends_on = None

label_mapping = {
    # Columbia
    1: ["C515",
        "D200",
        "D256",
        "D313",
        "C283",
        "D227",
        "D334",
        "D225",
        "D232",
        "D286",
        "D604",
        "F497",
        "D607",
        "F259",
        "F791",
        "D610",
        "E361",
        "D294",
        "C606",
        "D246",
        "C538",
        "C618",
        "D235",
        "D343",
        "F819",
        "F666",
        "H742",
        "D753",
        "D702",
        "D276",
        "D330",
        "D247",
        "D250",
        "D304",
        "D310",
        "C530",
        "D231",
        "D244",
        "D255",
        "D209",
        "D226",
        "D280",
        "D930",
        "D658",
        "D746",
        "E353",
        "E359",
        "F616",
        "D734",
        "F495",
        "C572",
        "D180",
        "F419",
        "D206",
        "D326",
        "D316",
        "C624",
        "D267",
        "C575",
        "D783",
        "F078",
        "F391",
        "E351",
        "F490",
        "F526",
        "D815",
        "D882",
        "D745",
        "F928",
        "H418",
        "H573",
        "K207",
        "C318",
        "C561",
        "D254",
        "C562",
        "D290",
        "D300",
        "D284",
        "C520",
        "C535",
        "C539",
        "D214",
        "E348",
        "E356",
        "D588",
        "D703",
        "E061",
        "E365",
        "F655",
        "F503",
        "M230",
        "C578",
        "D377",
        "C511",
        "C590",
        "C594",
        "D325",
        "D592",
        "E364",
        "D752",
        "E358",
        "E363",
        "F590",
        "F501",
        "E354",
        "E357",
        "D596",
        "D270",
        "D275",
        "D266",
        "M344",
        "C537",
        "D243",
        "D289",
        "D215",
        "D339",
        "D292",
        "D606",
        "I486",
        "E343",
        "E160",
        "F969",
        "K307",
        "F499",
        "E355",
        "F773",
        "C560",
        "C571",
        "C603",
        "D208",
        "C581",
        "D299",
        "E229",
        "F033",
        "F505",
        "H743",
        "E360",
        "F459",
        "H741",
        "D590",
        "E362"
        ],
    # Legacy
    4: [
        "D239",
        "D312",
        "C585",
        "D311",
        "D360",
        "D362",
        "D469",
        "D580",
        "E126",
        "E890",
        "F130",
        "F462",
        "D398",
        "F142",
        "D418",
        "D458",
        "D464",
        "E926",
        "E931",
        "D821",
        "D460",
        "D489",
        "D420",
        "J640",
        "C545",
        "C555",
        "C556",
        "C558",
        "C569",
        "D392",
        "D365",
        "C188",
        "C567",
        "D820",
        "D493",
        "D502",
        "E895",
        "E949",
        "D470",
        "D506",
        "D510",
        "D402",
        "D482",
        "D498",
        "D639",
        "E903",
        "F393",
        "D593",
        "K100",
        "M699",
        "C554",
        "C589",
        "C620",
        "D356",
        "D394",
        "C592",
        "D363",
        "C615",
        "D491",
        "E894",
        "D481",
        "E902",
        "E948",
        "F355",
        "F357",
        "D486",
        "E925",
        "F243",
        "D371",
        "D814",
        "D598",
        "C549",
        "C633",
        "D189",
        "D387",
        "A941",
        "C587",
        "D025",
        "D261",
        "D355",
        "C577",
        "C625",
        "D268",
        "D281",
        "D367",
        "C564",
        "D354",
        "D429",
        "D454",
        "F545",
        "D462",
        "E933",
        "H009",
        "D492",
        "D403",
        "D421",
        "D509",
        "E938",
        "E916",
        "D504",
        "D478",
        "H422",
        "M670",
        "D361",
        "D375",
        "D381",
        "C550",
        "D184",
        "D185",
        "D369",
        "D379",
        "D484",
        "D485",
        "F356",
        "F747",
        "D417",
        "D483",
        "D487",
        "D591",
        "F055",
        "F556",
        "E910",
        "D416",
        "D579",
        "D609",
        "E900",
        "F137",
        "C285",
        "D216",
        "C604",
        "C540",
        "C591",
        "D179",
        "D395",
        "B330",
        "C521",
        "C536",
        "D412",
        "D473",
        "F123",
        "F353",
        "D488",
        "D587",
        "D466",
        "D490",
        "D529",
        "V670",
        "C130",
        "C579",
        "D353",
        "C568",
        "D188",
        "D186",
        "D446",
        "E031",
        "D406",
        "D497",
        "D576",
        "D683",
        "D494",
        "G192",
        "D468",
        "D501",
        "C541",
        "C563",
        "B185",
        "C544",
        "D257",
        "D342",
        "C517",
        "D183",
        "E932",
        "D505",
        "D574",
        "E892",
        "D507",
        "D517",
        "E905",
        "E929",
        "D496",
        "H132"
    ],
    # CMG
    44: ['G102', 'G103', 'M229', 'M844']
}


def get_table(bind, table_name: str):
    meta = sa.MetaData(bind=bind)
    meta.reflect()
    return meta.tables[table_name]


def upgrade():
    op.add_column('Label', sa.Column('rep_owner_key', sa.String(), nullable=True))
    op.add_column('Label', sa.Column('child_rep_owner_keys', sa.ARRAY(sa.String()), nullable=True))
    op.execute(
        """
            UPDATE "Label"
            SET rep_owner_key = external_id,
                child_rep_owner_keys = ARRAY[external_id];
        """
    )
    op.drop_column('Label', 'external_id')
    op.create_unique_constraint('uq_Label_rep_owner_key', 'Label', ['rep_owner_key'])
    op.alter_column('Label', 'child_rep_owner_keys', existing_type=sa.ARRAY(sa.String()), nullable=False)
    for key, value in label_mapping.items():
        ss = "{ " + f"{', '.join(value)}" + " }"
        op.execute(
            f"""
                UPDATE "Label"
                SET child_rep_owner_keys ='{ss}'
                WHERE id = {key};
            """
        )


def downgrade():
    op.add_column('Label', sa.Column('external_id', sa.String(), nullable=True))
    op.execute(
        '''
            UPDATE "Label"
            SET external_id = rep_owner_key,
                child_rep_owner_keys = ARRAY[external_id];
        '''
    )
    op.create_unique_constraint('uq_Label_external_id', 'Label', ['external_id'])
    op.drop_column('Label', 'rep_owner_key')
    op.drop_column('Label', 'child_rep_owner_keys')