"""Edit mediaserver foreign keys

Revision ID: ba377c6bd1c2
Revises: 742a7e247e0c

"""

# revision identifiers, used by Alembic.
revision = 'ba377c6bd1c2'
down_revision = '742a7e247e0c'

from alembic import op
from sqlalchemy import sql
import sqlalchemy as sa


mediaserver_table = sql.table('mediaserver',
                              sql.column('id'),
                              sql.column('name'))

trunkfeatures_table = sql.table('trunkfeatures',
                                sql.column('mediaserverid'))

groupfeatures_table = sql.table('groupfeatures',
                                sql.column('mediaserverid'))

mediaserver_id_column = (mediaserver_table.c.id,)

all_mds_id = (sql.select(mediaserver_id_column)
              .where(mediaserver_table.c.name == 'all_mds')
              .limit(1))

mds_main_id = (sql.select(mediaserver_id_column)
               .where(mediaserver_table.c.name == 'default')
               .limit(1))


def _fill_sip_trunk_mediaserverid():
    fill_sip_trunk_query = (trunkfeatures_table
                            .update()
                            .values(mediaserverid=all_mds_id)
                            .where(trunkfeatures_table.c.mediaserverid == None))

    op.execute(fill_sip_trunk_query)


def _update_groupfeatures_mediaserverid():
    update_groupfeatures_query = (groupfeatures_table
                                  .update()
                                  .values(mediaserverid=mds_main_id)
                                  .where(groupfeatures_table.c.mediaserverid == 0))

    op.execute(update_groupfeatures_query)


def _add_groupfeatures_foreign_key():
    op.create_foreign_key('groupfeatures_mediaserverid_fkey', 'groupfeatures',
                          'mediaserver', ['mediaserverid'], ['id'],
                          ondelete='RESTRICT')


def upgrade():
    _fill_sip_trunk_mediaserverid()

    op.alter_column('trunkfeatures', 'mediaserverid', nullable=False)

    _update_groupfeatures_mediaserverid()

    _add_groupfeatures_foreign_key()

    op.alter_column('groupfeatures', 'mediaserverid', server_default=None)


def downgrade():
    op.drop_constraint(constraint_name="groupfeatures_mediaserverid_fkey", table_name="groupfeatures", type_="foreignkey")
    op.alter_column('trunkfeatures', 'mediaserverid', nullable=True)
    op.alter_column('groupfeatures', 'mediaserverid', server_default='0')

    clear_sip_trunk_query = (trunkfeatures_table
                             .update()
                             .values(mediaserverid=None)
                             .where(trunkfeatures_table.c.mediaserverid == all_mds_id))

    clear_groupfeatures_query = (groupfeatures_table
                                 .update()
                                 .values(mediaserverid=0)
                                 .where(groupfeatures_table.c.mediaserverid == mds_main_id))

    op.execute(clear_sip_trunk_query)
    op.execute(clear_groupfeatures_query)
