"""Introduce role permissions [5994f9ad0489].

Revision ID: 5994f9ad0489
Revises: 0.21.1
Create Date: 2022-10-25 23:52:25.935344

"""

import uuid

import sqlalchemy as sa
import sqlmodel
from alembic import op
from sqlalchemy import or_, select

from zenml.utils.time_utils import utc_now

# revision identifiers, used by Alembic.
revision = "5994f9ad0489"
down_revision = "0.21.1"
branch_labels = None
depends_on = None


def upgrade() -> None:
    """Upgrade database schema and/or data, creating a new revision."""
    # Create the rolepermissionschema table to track which permissions a given
    #  role grants
    op.create_table(
        "rolepermissionschema",
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("role_id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
        sa.ForeignKeyConstraint(
            ["role_id"],
            ["roleschema.id"],
        ),
        sa.PrimaryKeyConstraint("name", "role_id"),
    )
    # get metadata from current connection
    meta = sa.MetaData()

    # pass in tuple with tables we want to reflect, otherwise whole database
    #  will get reflected
    meta.reflect(
        bind=op.get_bind(),
        only=(
            "rolepermissionschema",
            "roleschema",
            "userroleassignmentschema",
            "teamroleassignmentschema",
            "userschema",
        ),
    )

    # In order to ensure unique names on roles delete potential admin/guest role
    #  that might have been created previous to this alembic version
    userroleassignmentschema = sa.Table(
        "userroleassignmentschema",
        meta,
    )
    teamroleassignmentschema = sa.Table(
        "teamroleassignmentschema",
        meta,
    )
    roleschema = sa.Table(
        "roleschema",
        meta,
    )
    conn = op.get_bind()
    res = conn.execute(
        select(roleschema.c.id).where(
            roleschema.c.name.in_(["admin", "guest"])
        )
    ).fetchall()

    role_ids = [i[0] for i in res]

    conn.execute(
        userroleassignmentschema.delete().where(
            userroleassignmentschema.c.role_id.in_(role_ids)
        )
    )
    conn.execute(
        teamroleassignmentschema.delete().where(
            teamroleassignmentschema.c.role_id.in_(role_ids)
        )
    )
    conn.execute(
        roleschema.delete().where(
            or_(roleschema.c.name == "admin", roleschema.c.name == "guest")
        )
    )

    # Create the three standard permissions also defined in
    #  zenml.enums.PermissionType
    read = "read"
    write = "write"
    me = "me"

    admin_id = str(uuid.uuid4()).replace("-", "")
    guest_id = str(uuid.uuid4()).replace("-", "")

    # Prefill the roles table with the admin and guest role
    now = utc_now()
    op.bulk_insert(
        sa.Table(
            "roleschema",
            meta,
        ),
        [
            {
                "id": admin_id,
                "name": "admin",
                "created": now,
                "updated": now,
            },
            {
                "id": guest_id,
                "name": "guest",
                "created": now,
                "updated": now,
            },
        ],
    )

    # Give the admin read, write and me permissions,
    # give the guest read and me permissions
    op.bulk_insert(
        sa.Table(
            "rolepermissionschema",
            meta,
        ),
        [
            {"role_id": admin_id, "name": read},
            {"role_id": admin_id, "name": write},
            {"role_id": admin_id, "name": me},
            {"role_id": guest_id, "name": read},
            {"role_id": guest_id, "name": me},
        ],
    )

    # In order to not break permissions for existing users, all existing users
    #  will be assigned the admin role
    userschema = sa.Table(
        "userschema",
        meta,
    )

    conn = op.get_bind()
    res = conn.execute(select(userschema.c.id)).fetchall()
    user_ids = [i[0] for i in res]

    now = utc_now()
    for user_id in user_ids:
        op.bulk_insert(
            sa.Table(
                "userroleassignmentschema",
                meta,
            ),
            [
                {
                    "id": str(uuid.uuid4()).replace("-", ""),
                    "role_id": admin_id,
                    "user_id": user_id,
                    "created": now,
                    "updated": now,
                }
            ],
        )


def downgrade() -> None:
    """Downgrade database schema and/or data back to the previous revision."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("rolepermissionschema")
    # ### end Alembic commands ###
