"""add_vfolder_type_column

Revision ID: 529113b08c2c
Revises: c481d3dc6c7d
Create Date: 2020-04-09 16:37:35.460936

"""
import textwrap
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
from sqlalchemy.sql.expression import bindparam
from ai.backend.manager.models import VFolderPermission, VFolderUsageMode, VFolderOwnershipType
from ai.backend.manager.models.base import convention, EnumValueType, IDColumn, GUID


# revision identifiers, used by Alembic.
revision = '529113b08c2c'
down_revision = 'c481d3dc6c7d'
branch_labels = None
depends_on = None

vfperm_choices = list(map(lambda v: v.value, VFolderPermission))
# vfolderpermission type should already be defined.

vfusagemode_choices = list(map(lambda v: v.value, VFolderUsageMode))
vfolderusagemode = postgresql.ENUM(
    *vfusagemode_choices,
    name='vfolderusagemode',
)

vfownershiptype_choices= list(map(lambda v: v.value, VFolderOwnershipType))
vfolderownershiptype = postgresql.ENUM(
    *vfownershiptype_choices,
    name='vfolderownershiptype',
)


def upgrade():
    metadata = sa.MetaData(naming_convention=convention)
    # partial table to be preserved and referred
    vfolders = sa.Table(
        'vfolders', metadata,
        IDColumn('id'),
        sa.Column('ownership_type', EnumValueType(VFolderOwnershipType), nullable=False),
        sa.Column('user', GUID, sa.ForeignKey('users.uuid'), nullable=True),
        sa.Column('group', GUID, sa.ForeignKey('groups.id'), nullable=True),
    )

    vfolderusagemode.create(op.get_bind())
    vfolderownershiptype.create(op.get_bind())
    op.add_column('vfolder_invitations',
                  sa.Column('modified_at', sa.DateTime(timezone=True), nullable=True,
                            onupdate=sa.func.current_timestamp()))
    op.add_column(
        'vfolders',
        sa.Column('usage_mode', sa.Enum(*vfusagemode_choices, name='vfolderusagemode'), nullable=True)
    )
    op.add_column(
        'vfolders',
        sa.Column('ownership_type',
                  sa.Enum(*vfownershiptype_choices, name='vfolderownershiptype'),
                  nullable=True)
    )
    op.add_column(
        'vfolders',
        sa.Column('permission',
                  sa.Enum(*vfperm_choices, name='vfolderpermission'),
                  nullable=True)
    )

    # Fill vfolders.c.usage_mode with 'general' and vfolders.c.permission.
    conn = op.get_bind()
    query = textwrap.dedent("UPDATE vfolders SET usage_mode = 'general';")
    conn.execute(query)
    query = textwrap.dedent("UPDATE vfolders SET permission = 'wd' WHERE \"user\" IS NOT NULL;")
    conn.execute(query)
    query = textwrap.dedent("UPDATE vfolders SET permission = 'rw' WHERE \"group\" IS NOT NULL;")
    conn.execute(query)

    # Set vfolders.c.ownership_type field based on user and group column.
    query = (sa.select([vfolders.c.id, vfolders.c.user, vfolders.c.group])
               .select_from(vfolders))
    updates = []
    for row in conn.execute(query).fetchall():
        if row['group']:
            ownership_type = VFolderOwnershipType.GROUP
        else:
            ownership_type = VFolderOwnershipType.USER
        updates.append({'vfid': row['id'], 'otype': ownership_type})
    if updates:
        query = (sa.update(vfolders)
                   .values(ownership_type=bindparam('otype'))
                   .where(vfolders.c.id == bindparam('vfid')))
        conn.execute(query, updates)

    # Create indexes for name.
    op.create_index(op.f('ix_vfolders_name'), 'vfolders', ['name'], unique=False)

    # Constraints
    op.create_check_constraint(
        'ownership_type_match_with_user_or_group',
        'vfolders',
        '(ownership_type = \'user\' AND "user" IS NOT NULL) OR '
        '(ownership_type = \'group\' AND "group" IS NOT NULL)'
    )
    op.create_check_constraint(
        'either_one_of_user_or_group',
        'vfolders',
        '("user" IS NULL AND "group" IS NOT NULL) OR ("user" IS NOT NULL AND "group" IS NULL)'
    )

    op.alter_column('vfolders', column_name='usage_mode', nullable=False)
    op.alter_column('vfolders', column_name='permission', nullable=False)


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint('ck_vfolders_ownership_type_match_with_user_or_group', 'vfolders')
    op.drop_constraint('ck_vfolders_either_one_of_user_or_group', 'vfolders')
    op.drop_index(op.f('ix_vfolders_name'), table_name='vfolders')
    op.drop_column('vfolders', 'usage_mode')
    op.drop_column('vfolders', 'ownership_type')
    op.drop_column('vfolders', 'permission')
    op.drop_column('vfolder_invitations', 'modified_at')
    vfolderusagemode.drop(op.get_bind())
    vfolderownershiptype.drop(op.get_bind())
    # ### end Alembic commands ###
