"""convert user.role to a mstr_role FK, restrict mstr_role.role_key to fixed tiers

Revision ID: e5f6a7b8c9d0
Revises: b9b3c25612c3
Create Date: 2026-07-23
"""
import uuid

from alembic import op
import sqlalchemy as sa

revision = 'e5f6a7b8c9d0'
down_revision = 'b9b3c25612c3'
branch_labels = None
depends_on = None

ROLE_KEY_HIERARCHY = ('SUPER_ADMIN', 'ADMIN', 'SR_CONSULTANT', 'CONSULTANT')
FALLBACK_ROLE_KEY = 'CONSULTANT'
CHECK_CONSTRAINT_NAME = 'ck_mstr_role_role_key_allowed'


def upgrade() -> None:
    conn = op.get_bind()

    op.add_column('user', sa.Column('role_id', sa.dialects.postgresql.UUID(as_uuid=True), nullable=True))
    op.drop_constraint('mstr_role_role_key_key', 'mstr_role', type_='unique')

    distinct_roles = conn.execute(sa.text('SELECT DISTINCT role FROM "user"')).scalars().all()
    for role_value in distinct_roles:
        existing_id = conn.execute(
            sa.text('SELECT id FROM mstr_role WHERE role_key = :role_key LIMIT 1'),
            {'role_key': role_value},
        ).scalar()

        if existing_id is None:
            new_id = uuid.uuid4()
            role_key = role_value if role_value in ROLE_KEY_HIERARCHY else FALLBACK_ROLE_KEY
            name = role_value.replace('_', ' ').title()
            # Role names are unique — a legacy value colliding with an existing name
            # falls back to a disambiguated one rather than failing the migration.
            name_taken = conn.execute(
                sa.text('SELECT 1 FROM mstr_role WHERE name = :name'), {'name': name}
            ).scalar()
            if name_taken:
                name = f'{name} ({role_value})'
            conn.execute(
                sa.text(
                    'INSERT INTO mstr_role (id, name, role_key, created_at) '
                    'VALUES (:id, :name, :role_key, now())'
                ),
                {'id': new_id, 'name': name, 'role_key': role_key},
            )
            existing_id = new_id

        conn.execute(
            sa.text('UPDATE "user" SET role_id = :role_id WHERE role = :role_value'),
            {'role_id': existing_id, 'role_value': role_value},
        )

    placeholders = ', '.join(f"'{k}'" for k in ROLE_KEY_HIERARCHY)
    conn.execute(sa.text(f"UPDATE mstr_role SET role_key = '{FALLBACK_ROLE_KEY}' WHERE role_key NOT IN ({placeholders})"))

    op.create_check_constraint(CHECK_CONSTRAINT_NAME, 'mstr_role', f"role_key IN ({placeholders})")

    op.alter_column('user', 'role_id', nullable=False)
    op.create_foreign_key('fk_user_role_id_mstr_role', 'user', 'mstr_role', ['role_id'], ['id'])
    op.drop_column('user', 'role')


def downgrade() -> None:
    op.add_column('user', sa.Column('role', sa.String(30), nullable=True))
    op.execute(
        'UPDATE "user" SET role = mstr_role.role_key FROM mstr_role WHERE "user".role_id = mstr_role.id'
    )
    op.alter_column('user', 'role', nullable=False)
    op.drop_constraint('fk_user_role_id_mstr_role', 'user', type_='foreignkey')
    op.drop_column('user', 'role_id')

    op.drop_constraint(CHECK_CONSTRAINT_NAME, 'mstr_role', type_='check')
    op.create_unique_constraint('mstr_role_role_key_key', 'mstr_role', ['role_key'])