"""Add curator fields

Revision ID: 351faebd379d
Revises: ee3f4b47fad5
Create Date: 2024-08-15 22:37:08.397052

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = "351faebd379d"
down_revision = "ee3f4b47fad5"
branch_labels: None = None
depends_on: None = None


def upgrade() -> None:
    # Add is_curator column to User__UserGroup table
    op.add_column(
        "user__user_group",
        sa.Column("is_curator", sa.Boolean(), nullable=False, server_default="false"),
    )

    # Use batch mode to modify the enum type
    with op.batch_alter_table("user", schema=None) as batch_op:
        batch_op.alter_column(  # type: ignore[attr-defined]
            "role",
            type_=sa.Enum(
                "BASIC",
                "ADMIN",
                "CURATOR",
                "GLOBAL_CURATOR",
                name="userrole",
                native_enum=False,
            ),
            existing_type=sa.Enum("BASIC", "ADMIN", name="userrole", native_enum=False),
            existing_nullable=False,
        )
    # Create the association table
    op.create_table(
        "credential__user_group",
        sa.Column("credential_id", sa.Integer(), nullable=False),
        sa.Column("user_group_id", sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(
            ["credential_id"],
            ["credential.id"],
        ),
        sa.ForeignKeyConstraint(
            ["user_group_id"],
            ["user_group.id"],
        ),
        sa.PrimaryKeyConstraint("credential_id", "user_group_id"),
    )
    op.add_column(
        "credential",
        sa.Column(
            "curator_public", sa.Boolean(), nullable=False, server_default="false"
        ),
    )


def downgrade() -> None:
    # Update existing records to ensure they fit within the BASIC/ADMIN roles
    op.execute(
        "UPDATE \"user\" SET role = 'ADMIN' WHERE role IN ('CURATOR', 'GLOBAL_CURATOR')"
    )

    # Remove is_curator column from User__UserGroup table
    op.drop_column("user__user_group", "is_curator")

    with op.batch_alter_table("user", schema=None) as batch_op:
        batch_op.alter_column(  # type: ignore[attr-defined]
            "role",
            type_=sa.Enum(
                "BASIC", "ADMIN", name="userrole", native_enum=False, length=20
            ),
            existing_type=sa.Enum(
                "BASIC",
                "ADMIN",
                "CURATOR",
                "GLOBAL_CURATOR",
                name="userrole",
                native_enum=False,
            ),
            existing_nullable=False,
        )
    # Drop the association table
    op.drop_table("credential__user_group")
    op.drop_column("credential", "curator_public")
