"""add_vfolder_host_permission

Revision ID: 5bce905c21e5
Revises: 360af8f33d4e
Create Date: 2022-10-13 18:35:21.955941

"""
from typing import Dict, List

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql as pgsql
from sqlalchemy.sql.expression import bindparam

from ai.backend.manager.models.base import GUID, VFolderHostPermission, metadata

# revision identifiers, used by Alembic.
revision = "5bce905c21e5"
down_revision = "360af8f33d4e"
branch_labels = None
depends_on = None


ALL_HOST_PERMISSIONS = [perm.value for perm in VFolderHostPermission]


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    domains = sa.Table(
        "domains",
        metadata,
        sa.Column("name", sa.String(length=64), primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.ARRAY(sa.String)),
        extend_existing=True,
    )
    groups = sa.Table(
        "groups",
        metadata,
        sa.Column("id", GUID, primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.ARRAY(sa.String)),
        extend_existing=True,
    )
    keypair_resource_policies = sa.Table(
        "keypair_resource_policies",
        metadata,
        sa.Column("name", sa.String(length=256), primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.ARRAY(sa.String)),
        extend_existing=True,
    )
    domain_query = sa.select([domains.c.name, domains.c.allowed_vfolder_hosts])
    group_query = sa.select([groups.c.id, groups.c.allowed_vfolder_hosts])
    keypair_resource_policies_query = sa.select(
        [keypair_resource_policies.c.name, keypair_resource_policies.c.allowed_vfolder_hosts]
    )

    connection = op.get_bind()
    bind_param_pk = "bind_param_pk"

    def get_pk_vfolder_host_maps(query, pk_name) -> List:
        map_list = []
        rows = connection.execute(query).fetchall()
        for row in rows:
            hosts: List[str] = row["allowed_vfolder_hosts"]
            all_perms: Dict[str, List[str]] = {host: ALL_HOST_PERMISSIONS for host in hosts}
            map_list.append({bind_param_pk: row[pk_name], "allowed_vfolder_hosts": all_perms})
        return map_list

    domain_list = get_pk_vfolder_host_maps(domain_query, "name")
    group_list = get_pk_vfolder_host_maps(group_query, "id")
    resource_policy_list = get_pk_vfolder_host_maps(keypair_resource_policies_query, "name")

    op.drop_column("domains", "allowed_vfolder_hosts")
    op.drop_column("groups", "allowed_vfolder_hosts")
    op.drop_column("keypair_resource_policies", "allowed_vfolder_hosts")

    op.add_column(
        "domains",
        sa.Column(
            "allowed_vfolder_hosts",
            pgsql.JSONB(),
            nullable=False,
            server_default=sa.text("'{}'::jsonb"),
        ),
    )
    op.add_column(
        "groups",
        sa.Column(
            "allowed_vfolder_hosts",
            pgsql.JSONB(),
            nullable=False,
            server_default=sa.text("'{}'::jsonb"),
        ),
    )
    op.add_column(
        "keypair_resource_policies",
        sa.Column(
            "allowed_vfolder_hosts",
            pgsql.JSONB(),
            nullable=False,
            server_default=sa.text("'{}'::jsonb"),
        ),
    )

    domains = sa.Table(
        "domains",
        metadata,
        sa.Column("name", sa.String(length=64), primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.JSONB),
        extend_existing=True,
    )
    groups = sa.Table(
        "groups",
        metadata,
        sa.Column("id", GUID, primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.JSONB),
        extend_existing=True,
    )
    keypair_resource_policies = sa.Table(
        "keypair_resource_policies",
        metadata,
        sa.Column("name", sa.String(length=256), primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.JSONB),
        extend_existing=True,
    )

    def update_vfolder_host_col(table, cond, row_list):
        query = (
            sa.update(table)
            .values({"allowed_vfolder_hosts": bindparam("allowed_vfolder_hosts")})
            .where(cond)
        )
        connection.execute(query, row_list)

    update_vfolder_host_col(domains, domains.c.name == bindparam(bind_param_pk), domain_list)
    update_vfolder_host_col(groups, groups.c.id == bindparam(bind_param_pk), group_list)
    update_vfolder_host_col(
        keypair_resource_policies,
        keypair_resource_policies.c.name == bindparam(bind_param_pk),
        resource_policy_list,
    )

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    domains = sa.Table(
        "domains",
        metadata,
        sa.Column("name", sa.String(length=64), primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.JSONB),
        extend_existing=True,
    )
    groups = sa.Table(
        "groups",
        metadata,
        sa.Column("id", GUID, primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.JSONB),
        extend_existing=True,
    )
    keypair_resource_policies = sa.Table(
        "keypair_resource_policies",
        metadata,
        sa.Column("name", sa.String(length=256), primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.JSONB),
        extend_existing=True,
    )
    connection = op.get_bind()
    bind_param_pk = "bind_param_pk"

    def get_pk_vfolder_host_maps(query, pk_name) -> List:
        map_list = []
        rows = connection.execute(query).fetchall()
        for row in rows:
            hosts: Dict[str, List[str]] = row["allowed_vfolder_hosts"]
            map_list.append(
                {bind_param_pk: row[pk_name], "allowed_vfolder_hosts": list(hosts.keys())}
            )
        return map_list

    domain_query = sa.select([domains.c.name, domains.c.allowed_vfolder_hosts])
    group_query = sa.select([groups.c.id, groups.c.allowed_vfolder_hosts])
    keypair_resource_policies_query = sa.select(
        [keypair_resource_policies.c.name, keypair_resource_policies.c.allowed_vfolder_hosts]
    )

    domain_list = get_pk_vfolder_host_maps(domain_query, "name")
    group_list = get_pk_vfolder_host_maps(group_query, "id")
    resource_policy_list = get_pk_vfolder_host_maps(keypair_resource_policies_query, "name")

    op.drop_column("domains", "allowed_vfolder_hosts")
    op.drop_column("groups", "allowed_vfolder_hosts")
    op.drop_column("keypair_resource_policies", "allowed_vfolder_hosts")

    op.add_column(
        "domains",
        sa.Column(
            "allowed_vfolder_hosts", pgsql.ARRAY(sa.String), nullable=False, server_default="{}"
        ),
    )
    op.add_column(
        "groups",
        sa.Column(
            "allowed_vfolder_hosts", pgsql.ARRAY(sa.String), nullable=False, server_default="{}"
        ),
    )
    op.add_column(
        "keypair_resource_policies",
        sa.Column(
            "allowed_vfolder_hosts", pgsql.ARRAY(sa.String), nullable=False, server_default="{}"
        ),
    )

    domains = sa.Table(
        "domains",
        metadata,
        sa.Column("name", sa.String(length=64), primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.ARRAY(sa.String)),
        extend_existing=True,
    )
    groups = sa.Table(
        "groups",
        metadata,
        sa.Column("id", GUID, primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.ARRAY(sa.String)),
        extend_existing=True,
    )
    keypair_resource_policies = sa.Table(
        "keypair_resource_policies",
        metadata,
        sa.Column("name", sa.String(length=256), primary_key=True),
        sa.Column("allowed_vfolder_hosts", pgsql.ARRAY(sa.String)),
        extend_existing=True,
    )

    def update_vfolder_host_col(table, cond, row_list):
        query = (
            sa.update(table)
            .values({"allowed_vfolder_hosts": bindparam("allowed_vfolder_hosts")})
            .where(cond)
        )
        connection.execute(query, row_list)

    update_vfolder_host_col(domains, domains.c.name == bindparam(bind_param_pk), domain_list)
    update_vfolder_host_col(groups, groups.c.id == bindparam(bind_param_pk), group_list)
    update_vfolder_host_col(
        keypair_resource_policies,
        keypair_resource_policies.c.name == bindparam(bind_param_pk),
        resource_policy_list,
    )

    # ### end Alembic commands ###
