"""add_country_codes_to_territories

Revision ID: dc82a739de05
Revises: cd934930f352
Create Date: 2020-09-11 06:47:18.043764

"""
from alembic import op
import sqlalchemy as sa
from migrations.fixtures.country_codes import country_codes

# revision identifiers, used by Alembic.
revision = 'dc82a739de05'
down_revision = 'd34838b29791'
branch_labels = None
depends_on = None


def upgrade():
    bind = op.get_bind()
    op.add_column('Territory', sa.Column('code', sa.String(), nullable=True))

    bind.execute(f''' INSERT INTO "TerritoryType" (id, name)
                        VALUES (0, 'Worldwide'), (1, 'Country') ON CONFLICT DO NOTHING
                        ''')
    bind.execute(f''' INSERT INTO "Territory" (id, name, code, territory_type_id, parent_id)
                        VALUES (0, 'Global', 'GLBL', 0, Null) ON CONFLICT DO NOTHING
                        ''')

    bind.execute('''ALTER SEQUENCE "Territory_id_seq" RESTART WITH 248''')

    bind.execute(f''' UPDATE "Territory" SET code='GLBL' WHERE id = 0''')

    for code, name in country_codes.items():
        territory_id = bind.execute(f''' SELECT id FROM "Territory" WHERE name='{name}' ''').scalar()
        if territory_id:
            bind.execute(f''' UPDATE "Territory" SET code='{code}' WHERE id = {territory_id}''')
        else:
            bind.execute(f''' INSERT INTO "Territory" (name, code, territory_type_id, parent_id)
                              VALUES ('{name}', '{code}', 1, 0) ''')

    bind.execute(f''' INSERT INTO "Territory" (name, code, territory_type_id, parent_id)
                                  VALUES ('Democratic Republic of the Congo', 'CD', 1, 0) ''')

    op.alter_column('Territory', 'code', nullable=False)


def downgrade():
    bind = op.get_bind()
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('Territory', 'code')
    bind.execute(f''' DELETE FROM "Territory" WHERE name = 'Democratic Republic of the Congo' ''')
    # ### end Alembic commands ###