"""remove stale startup permissions

Revision ID: 20260527_04
Revises: 20260527_03
Create Date: 2026-05-27 11:05:00.000000

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


revision: str = "20260527_04"
down_revision: Union[str, None] = "20260527_03"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


STALE_STARTUP_PERMISSION_KEYS = (
    "budget:create",
    "budget:list",
    "budget:read",
    "budget:update",
    "budget:delete",
    "timeline:create",
    "timeline:list",
    "timeline:read",
    "timeline:update",
    "timeline:delete",
)


def _quoted_keys() -> str:
    return ", ".join(f"'{key}'" for key in STALE_STARTUP_PERMISSION_KEYS)


def _table_exists(inspector: sa.Inspector, table_name: str) -> bool:
    return table_name in inspector.get_table_names()


def _remove_template_permission_keys(table_name: str, *, touch_updated_at: bool = False) -> None:
    updated_at_assignment = ", updated_at = NOW()" if touch_updated_at else ""
    op.execute(
        f"""
        UPDATE {table_name}
        SET permissions = (
            SELECT jsonb_object_agg(role_key, role_permissions - ARRAY[{_quoted_keys()}])
            FROM jsonb_each(permissions::jsonb) AS role_entries(role_key, role_permissions)
        ){updated_at_assignment}
        WHERE EXISTS (
            SELECT 1
            FROM jsonb_each(permissions::jsonb) AS role_entries(role_key, role_permissions)
            WHERE role_permissions ?| ARRAY[{_quoted_keys()}]
        )
        """
    )


def upgrade() -> None:
    bind = op.get_bind()
    inspector = sa.inspect(bind)

    if _table_exists(inspector, "api_endpoint_permissions"):
        op.execute(
            f"""
            DELETE FROM api_endpoint_permissions
            WHERE endpoint_key IN ({_quoted_keys()})
            """
        )

    if _table_exists(inspector, "permission_templates"):
        _remove_template_permission_keys("permission_templates", touch_updated_at=True)

    if _table_exists(inspector, "permission_template_versions"):
        _remove_template_permission_keys("permission_template_versions")


def downgrade() -> None:
    pass