"""add_vfolder_host_permission

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

"""

from typing import Any

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

from ai.backend.common.types import VFolderHostPermission
from ai.backend.manager.models.base import GUID, 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: sa.sql.Select[Any], pk_name: str
    ) -> list[dict[str, object]]:
        map_list = []
        rows = connection.execute(query).fetchall()
        for row in rows:
            hosts: list[str] = row.allowed_vfolder_hosts
            all_perms: dict[str, list[str]] = dict.fromkeys(hosts, ALL_HOST_PERMISSIONS)
            map_list.append({
                bind_param_pk: getattr(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: sa.Table, cond: sa.ColumnElement[Any], row_list: list[dict[str, object]]
    ) -> None:
        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: sa.sql.Select[Any], pk_name: str
    ) -> list[dict[str, object]]:
        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: getattr(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: sa.Table, cond: sa.ColumnElement[Any], row_list: list[dict[str, object]]
    ) -> None:
        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 ###
