"""Enable RLS on critical tables missing it.

Tables: api_tokens, sequences, sessions, tenant_plugin_activation,
user_tenants, password_reset_tokens.
Also adds tenant_id to guest_invitations and enables RLS.

Revision ID: 0101
"""

from alembic import op
import sqlalchemy as sa

revision = "0101"
down_revision = "0100"
branch_labels = None
depends_on = None

# Tables that have tenant_id but no RLS
RLS_TABLES = [
    "api_tokens",
    "sequences",
    "sessions",
    "tenant_plugin_activation",
    "user_tenants",
    "password_reset_tokens",
]


def upgrade() -> None:
    # Enable RLS + create tenant isolation policy for each table
    for table in RLS_TABLES:
        op.execute(f"ALTER TABLE {table} ENABLE ROW LEVEL SECURITY;")
        op.execute(
            f"CREATE POLICY {table}_tenant_isolation ON {table} "
            f"FOR ALL USING (tenant_id = current_setting('app.current_tenant_id')::uuid) "
            f"WITH CHECK (tenant_id = current_setting('app.current_tenant_id')::uuid);"
        )

    # guest_invitations: add tenant_id + enable RLS
    op.add_column("guest_invitations", sa.Column("tenant_id", sa.UUID(), nullable=True))
    op.execute("CREATE INDEX ix_guest_invitations_tenant_id ON guest_invitations (tenant_id);")
    op.execute("ALTER TABLE guest_invitations ENABLE ROW LEVEL SECURITY;")
    op.execute(
        "CREATE POLICY guest_invitations_tenant_isolation ON guest_invitations "
        "FOR ALL USING (tenant_id = current_setting('app.current_tenant_id')::uuid) "
        "WITH CHECK (tenant_id = current_setting('app.current_tenant_id')::uuid);"
    )


def downgrade() -> None:
    # Drop guest_invitations RLS + tenant_id
    op.execute("DROP POLICY IF EXISTS guest_invitations_tenant_isolation ON guest_invitations;")
    op.execute("ALTER TABLE guest_invitations DISABLE ROW LEVEL SECURITY;")
    op.execute("DROP INDEX IF EXISTS ix_guest_invitations_tenant_id;")
    op.drop_column("guest_invitations", "tenant_id")

    # Drop RLS on other tables
    for table in RLS_TABLES:
        op.execute(f"DROP POLICY IF EXISTS {table}_tenant_isolation ON {table};")
        op.execute(f"ALTER TABLE {table} DISABLE ROW LEVEL SECURITY;")