from typing import List, Optional, Callable, Any, Tuple
import os
from pydantic import BaseModel
import casbin
from casbin_sqlalchemy_adapter import Adapter
# ← REMOVED: from casbin_redis_watcher import RedisWatcher, WatcherOptions
from sqlalchemy import and_, or_
from sqlalchemy.orm import Session, joinedload
from ..schema import CreatePermissionSchema
from ..models import (
    Role,
    RolePermission,
    UserRole,
    Permission,
    UserPermission
)
from abs_utils.logger import setup_logger

logger = setup_logger(__name__)
from abs_exception_core.exceptions import (
    DuplicatedError,
    NotFoundError,
    PermissionDeniedError
)

from ..models.gov_casbin_rule import GovCasbinRule
# ← REMOVED: from redis import Redis
# ← REMOVED: import time
import time

# ← NEW: Import the Service Bus watcher
from .watcher import AzureServiceBusWatcher, AzureServiceBusWatcherSchema


# ← REMOVED: RedisWatcherSchema class entirely
# The AzureServiceBusWatcherSchema from watcher.py replaces it


class RBACService:
    def __init__(self, session: Callable[..., Session], servicebus_config: Optional[AzureServiceBusWatcherSchema] = None):
        #                                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        #                                                ← CHANGED: was redis_config: Optional[RedisWatcherSchema] = None
        """
        Service For Managing the RBAC
        Args:
            session: Callable[...,Session] -> Session of the SQLAlchemy database engine
            servicebus_config: Optional config for Azure Service Bus watcher
        """
        self.db = session
        self.enforcer = None
        self.watcher = None
        self._initialize_casbin(servicebus_config)  # ← CHANGED: was redis_config


    def _save_policy_if_watcher_active(self):
        """
        Helper method to save policy only if watcher is active.
        """
        if self.is_watcher_active():
            self.enforcer.save_policy()

    def _bulk_operation_context(self):
        """
        Context manager for bulk Casbin operations.
        """
        class BulkOperationContext:
            def __init__(self, service):
                self.service = service

            def __enter__(self):
                return self

            def __exit__(self, exc_type, exc_val, exc_tb):
                if exc_type is None and self.service.is_watcher_active():
                    subscription = self.service.watcher.subscription_name if self.service.watcher else "unknown"
                    logger.info(f"[RBAC] 📤 Publishing policy change from subscription='{subscription}'...")
                    result = self.service.watcher.update()
                    if result:
                        logger.info(f"[RBAC] ✅ Policy change PUBLISHED from subscription='{subscription}'")
                    else:
                        logger.error(f"[RBAC] ❌ Policy change PUBLISH FAILED from subscription='{subscription}'")
                elif exc_type is None:
                    logger.warning(f"[RBAC] ⚠️ Watcher NOT active - policy change NOT published to other servers")
                return False

        return BulkOperationContext(self)

    def load_policy_with_retry(self, max_retries=3, delay=2):
        """Wraps the standard load_policy with retry logic. ← NO CHANGES"""
        for attempt in range(max_retries):
            try:
                self.enforcer.load_policy()
                logger.info("Casbin Policies loaded successfully.")
                break
            except Exception as e:
                logger.error(f"Error loading Casbin Policies: {e}")
                if attempt < max_retries - 1:
                    logger.warning(f"Error loading Casbin Policies (attempt {attempt + 1}). Retrying in {delay}s... Error: {e}")
                    time.sleep(delay)
                else:
                    logger.error("Max retries reached. Error loading Casbin Policies.")
                    raise

    def _on_policy_change(self, msg):
        """Callback when another server publishes a policy change."""
        subscription = self.watcher.subscription_name if self.watcher else "unknown"
        logger.info(f"[RBAC] 📩 Policy change received on subscription='{subscription}'")
        logger.info(f"[RBAC] 📩 Message: {msg}")

        policy_count_before = len(self.enforcer.get_policy())
        logger.info(f"[RBAC] 🔄 Reloading policies... (current count: {policy_count_before})")

        self.load_policy_with_retry()

        policy_count_after = len(self.enforcer.get_policy())
        logger.info(
            f"[RBAC] ✅ Policies reloaded on subscription='{subscription}': "
            f"{policy_count_before} → {policy_count_after} policies"
        )

    def _initialize_casbin(self, servicebus_config: Optional[AzureServiceBusWatcherSchema] = None):
        #                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        #                       ← CHANGED: was redis_config: Optional[RedisWatcherSchema] = None
        """
        Initiates the casbin policy using the default rules.
        ← CHANGED: Redis watcher replaced with Azure Service Bus watcher.
        """
        with self.db() as session:
            engine = session.get_bind()

            adapter = Adapter(engine, db_class=GovCasbinRule)
            
            current_dir = os.path.dirname(os.path.abspath(__file__))
            policy_path = os.path.join(current_dir, "policy.conf")
            
            self.enforcer = casbin.Enforcer(
                policy_path, adapter
            )
            self.enforcer.enable_auto_save(True)

            # ← CHANGED: Entire block below replaces the Redis watcher setup
            if servicebus_config:
                try:
                    watcher = AzureServiceBusWatcher(servicebus_config)
                    watcher.set_update_callback(
                        lambda msg: self._on_policy_change(msg)
                    )
                    watcher.start()

                    self.enforcer.set_watcher(watcher)
                    self.watcher = watcher
                    logger.info("Azure Service Bus watcher initialized successfully")

                except Exception as e:
                    logger.error(f"Failed to initialize Service Bus watcher: {e}")
                    self.watcher = None
            else:
                logger.info("Service Bus watcher not configured - Casbin will work without real-time policy updates")

    # =========================================================================
    # ALL METHODS BELOW ARE 100% UNCHANGED FROM YOUR EXISTING CODE
    # =========================================================================

    def add_policy(self, role: str, resource: str, action: str, module: str):
        with self._bulk_operation_context():
            self.enforcer.add_policy(role, resource, action, module)

    def remove_policy(self, role: str, resource: str, action: str, module: str):
        with self._bulk_operation_context():
            self.enforcer.remove_policy(role, resource, action, module)

    def add_policies(self, policies: List[Tuple[str, str, str, str]]):
        with self._bulk_operation_context():
            self.enforcer.add_policies(policies)

    def remove_policies(self, policies: List[List[str]]):
        with self._bulk_operation_context():
            self.enforcer.remove_policies(policies)

    def enforce_policy(self, role: str, resource: str, action: str, module: str):
        return self.enforcer.enforce(role, resource, action, module)
    
    def remove_filter_policy(self, index: int, value: str):
        with self._bulk_operation_context():
            self.enforcer.remove_filtered_policy(index, value)

    async def bulk_create_permissions(self, permissions: List[CreatePermissionSchema]):
        with self.db() as session:
            try:
                if not permissions:
                    return []
                if hasattr(permissions[0], 'model_dump'):
                    add_permissions = [Permission(**permission.model_dump()) for permission in permissions]
                else:
                    add_permissions = [Permission(**permission) for permission in permissions]
                session.bulk_save_objects(add_permissions)
                session.commit()
                return add_permissions
            except Exception:
                raise
            
    def build_filter(self, cond: dict):
        if "and" in cond:
            return and_(*[self.build_filter(c) for c in cond["and"]])
        elif "or" in cond:
            return or_(*[self.build_filter(c) for c in cond["or"]])
        else:
            return and_(*[
                getattr(Permission, field) == value
                for field, value in cond.items()
            ])

    async def get_permissions_by_condition(self, condition: dict):
        with self.db() as session:
            try:
                query = session.query(Permission).filter(self.build_filter(condition))
                return query.all()
            except Exception:
                raise

    async def delete_permission_by_uuids(self, permission_uuids: List[str]):
        with self.db() as session:
            try:
                session.query(UserPermission).filter(UserPermission.permission_uuid.in_(permission_uuids)).delete(synchronize_session=False)
                session.query(RolePermission).filter(RolePermission.permission_uuid.in_(permission_uuids)).delete(synchronize_session=False)
                permissions = session.query(Permission).filter(Permission.uuid.in_(permission_uuids))
                with self._bulk_operation_context():
                    for permission in permissions:
                        self.enforcer.remove_filtered_policy(1, permission.resource)
                permissions.delete(synchronize_session=False)
                session.commit()
                return True
            except Exception:
                raise

    def assign_permissions_to_user(self, user_uuid: str, permission_uuids: List[str]):
        with self.db() as session:
            try:
                current_permissions = session.query(UserPermission).filter(UserPermission.user_uuid == user_uuid).all()
                current_permission_uuids = [permission.permission_uuid for permission in current_permissions]
                remove_permissions = set(current_permission_uuids) - set(permission_uuids)
                add_permissions = set(permission_uuids) - set(current_permission_uuids)
                if remove_permissions:
                    self.revoke_user_permissions(user_uuid, list(remove_permissions))
                if add_permissions:
                    self.attach_permissions_to_user(user_uuid, list(add_permissions))
                return self.get_user_only_permissions(user_uuid)
            except Exception:
                raise

    def attach_permissions_to_user(self, user_uuid: str, permission_uuids: List[str]):
        with self.db() as session:
            try:
                user_permissions_data = [
                    {"user_uuid": user_uuid, "permission_uuid": permission_uuid}
                    for permission_uuid in permission_uuids
                ]
                session.bulk_insert_mappings(UserPermission, user_permissions_data)
                session.commit()
                permissions = session.query(Permission).filter(Permission.uuid.in_(permission_uuids)).all()
                policies = [
                    [f"user:{user_uuid}", permission.resource, permission.action, permission.module]
                    for permission in permissions
                ]
                with self._bulk_operation_context():
                    self.enforcer.add_policies(policies)
                return self.get_user_only_permissions(user_uuid)
            except Exception:
                raise

    def revoke_user_permissions(self, user_uuid: str, permission_uuids: List[str]):
        with self.db() as session:
            try:
                user_permissions = session.query(UserPermission).filter(
                    UserPermission.user_uuid == user_uuid,
                    UserPermission.permission_uuid.in_(permission_uuids)
                )
                permissions = session.query(Permission).filter(Permission.uuid.in_(permission_uuids)).all()
                policies = [
                    [f"user:{user_uuid}", permission.resource, permission.action, permission.module]
                    for permission in permissions
                ]
                with self._bulk_operation_context():
                    self.enforcer.remove_policies(policies)
                user_permissions.delete(synchronize_session=False)
                session.commit()
                return self.get_user_only_permissions(user_uuid)
            except Exception:
                raise

    def revoke_all_user_access(self, user_uuid: str) -> dict:
        with self.db() as session:
            try:
                if not session.is_active:
                    session.begin()
                user_roles = (
                    session.query(UserRole)
                    .options(joinedload(UserRole.role))
                    .filter(UserRole.user_uuid == user_uuid)
                    .all()
                )
                role_uuids = [ur.role.uuid for ur in user_roles] if user_roles else []
                user_permissions = (
                    session.query(UserPermission)
                    .options(joinedload(UserPermission.permission))
                    .filter(UserPermission.user_uuid == user_uuid)
                    .all()
                )
                permission_uuids = [up.permission.uuid for up in user_permissions] if user_permissions else []
                if role_uuids:
                    session.query(UserRole).filter(
                        UserRole.user_uuid == user_uuid,
                        UserRole.role_uuid.in_(role_uuids)
                    ).delete(synchronize_session=False)
                if permission_uuids:
                    permissions = session.query(Permission).filter(
                        Permission.uuid.in_(permission_uuids)
                    ).all()
                    policies = [
                        [f"user:{user_uuid}", permission.resource, permission.action, permission.module]
                        for permission in permissions
                    ]
                    with self._bulk_operation_context():
                        self.enforcer.remove_policies(policies)
                    session.query(UserPermission).filter(
                        UserPermission.user_uuid == user_uuid,
                        UserPermission.permission_uuid.in_(permission_uuids)
                    ).delete(synchronize_session=False)
                session.commit()
                return {
                    "user_uuid": user_uuid,
                    "roles_revoked": len(role_uuids),
                    "permissions_revoked": len(permission_uuids),
                    "role_uuids": role_uuids,
                    "permission_uuids": permission_uuids
                }
            except Exception:
                raise

    def list_roles(self) -> Any:
        with self.db() as session:
            try:
                total = session.query(Role).count()
                roles = session.query(Role).all()
                return {"roles": roles, "total": total}
            except Exception:
                raise

    def create_role(
        self,
        name: str,
        description: Optional[str] = None,
        permission_ids: List[str] = None,
    ) -> Any:
        with self.db() as session:
            try:
                existing_role = session.query(Role).filter(Role.name == name).first()
                if existing_role:
                    raise DuplicatedError(detail="Role already exists")
                role = Role(name=name, description=description)
                session.add(role)
                session.flush()
                if permission_ids:
                    permission_count = (
                        session.query(Permission)
                        .filter(Permission.uuid.in_(permission_ids))
                        .count()
                    )
                    if permission_count != len(permission_ids):
                        existing_permissions = (
                            session.query(Permission)
                            .filter(Permission.uuid.in_(permission_ids))
                            .all()
                        )
                        found_permission_ids = {p.uuid for p in existing_permissions}
                        missing_ids = set(permission_ids) - found_permission_ids
                        raise NotFoundError(
                            detail=f"Permissions with UUIDs '{', '.join(missing_ids)}' not found"
                        )
                    existing_permissions = (
                        session.query(Permission)
                        .filter(Permission.uuid.in_(permission_ids))
                        .all()
                    )
                    role_permissions = [
                        {"role_uuid": role.uuid, "permission_uuid": permission_uuid}
                        for permission_uuid in permission_ids
                    ]
                    session.bulk_insert_mappings(RolePermission, role_permissions)
                    policies = [
                        [role.uuid, permission.resource, permission.action, permission.module]
                        for permission in existing_permissions
                    ]
                    with self._bulk_operation_context():
                        self.enforcer.add_policies(policies)
                session.commit()
                session.refresh(role)
                return role
            except Exception:
                raise

    def get_role_with_permissions(self, role_uuid: str) -> Any:
        with self.db() as session:
            role = (
                session.query(Role)
                .options(joinedload(Role.permissions))
                .filter(Role.uuid == role_uuid)
                .first()
            )
            if not role:
                raise NotFoundError(detail="Requested role does not exist")
            return role

    def update_role_permissions(
        self,
        role_uuid: str,
        permissions: Optional[List[str]] = None,
        name: Optional[str] = None,
        description: Optional[str] = None,
    ) -> Any:
        with self.db() as session:
            try:
                if not session.is_active:
                    session.begin()
                role = (
                    session.query(Role)
                    .options(joinedload(Role.permissions))
                    .filter(Role.uuid == role_uuid)
                    .first()
                )
                if not role:
                    raise NotFoundError(detail="Requested role does not exist")
                if name is not None or description is not None:
                    if name:
                        existing_role = (
                            session.query(Role)
                            .filter(Role.name == name, Role.uuid != role_uuid)
                            .first()
                        )
                        if existing_role:
                            raise DuplicatedError(detail="Role already exists")
                        if role.name != "super_admin":
                            role.name = name
                    if description is not None:
                        role.description = description
                if permissions is not None:
                    with self._bulk_operation_context():
                        self.enforcer.remove_filtered_policy(0, str(role_uuid))
                        session.query(RolePermission).filter(
                            RolePermission.role_uuid == role_uuid
                        ).delete(synchronize_session=False)
                        if permissions:
                            permissions_objs = (
                                session.query(Permission)
                                .filter(Permission.uuid.in_(permissions))
                                .all()
                            )
                            found_permission_ids = {p.uuid for p in permissions_objs}
                            missing_permission_ids = set(permissions) - found_permission_ids
                            if missing_permission_ids:
                                raise NotFoundError(
                                    detail=f"Permissions with UUIDs '{', '.join(missing_permission_ids)}' not found"
                                )
                            role_permissions = [
                                {"role_uuid": role_uuid, "permission_uuid": permission.uuid}
                                for permission in permissions_objs
                            ]
                            session.bulk_insert_mappings(RolePermission, role_permissions)
                            policies = [
                                [role_uuid, permission.resource, permission.action, permission.module]
                                for permission in permissions_objs
                            ]
                            self.enforcer.add_policies(policies)
                session.commit()
                session.refresh(role)
                return role
            except Exception:
                raise

    def delete_role(self, role_uuid: str, exception_roles: List[str] = None):
        with self.db() as session:
            try:
                if not session.is_active:
                    session.begin()
                role = (
                    session.query(Role)
                    .options(joinedload(Role.permissions))
                    .filter(Role.uuid == role_uuid)
                    .first()
                )
                if not role:
                    raise NotFoundError(detail="Requested role does not exist")
                if exception_roles and len(exception_roles) > 0 and role.name in exception_roles:
                    raise PermissionDeniedError(detail="You are not allowed to delete the requested role.")
                remove_policies = [
                    [role.uuid, permission.resource, permission.action, permission.module]
                    for permission in role.permissions
                ]
                if remove_policies:
                    with self._bulk_operation_context():
                        self.enforcer.remove_policies(remove_policies)
                session.delete(role)
                session.commit()
            except Exception:
                raise

    def list_permissions(self) -> List[Any]:
        with self.db() as session:
            return session.query(Permission).all()

    def list_module_permissions(self, module: str) -> List[Any]:
        with self.db() as session:
            return session.query(Permission).filter(Permission.module == module).all()

    def get_user_only_permissions(self, user_uuid: str) -> List[Any]:
        with self.db() as session:
            user_permissions = (
                session.query(UserPermission)
                .filter(UserPermission.user_uuid == user_uuid)
                .options(joinedload(UserPermission.permission))
                .all()
            )
            result = []
            for user_permission in user_permissions:
                result.append(
                    {
                        "permission_id": user_permission.permission.uuid,
                        "created_at": user_permission.permission.created_at,
                        "updated_at": user_permission.permission.updated_at,
                        "name": user_permission.permission.name,
                        "resource": user_permission.permission.resource,
                        "action": user_permission.permission.action,
                        "module": user_permission.permission.module
                    }
                )
            return result

    def get_user_permissions(self, user_uuid: str) -> List[Any]:
        with self.db() as session:
            user_roles = (
                session.query(UserRole)
                .join(Role, UserRole.role_uuid == Role.uuid)
                .options(
                    joinedload(UserRole.role).joinedload(Role.permissions)
                )
                .filter(UserRole.user_uuid == user_uuid)
                .all()
            )
            if not user_roles:
                return []
            result = []
            for user_role in user_roles:
                role = user_role.role
                for permission in role.permissions:
                    result.append(
                        {
                            "permission_id": permission.uuid,
                            "created_at": permission.created_at,
                            "role_id": role.uuid,
                            "updated_at": permission.updated_at,
                            "role_name": role.name,
                            "name": permission.name,
                            "resource": permission.resource,
                            "action": permission.action,
                            "module": permission.module
                        }
                    )
            return result

    def bulk_revoke_permissions(self, role_uuid: str, permission_uuids: List[str]) -> Any:
        with self.db() as session:
            try:
                if not session.is_active:
                    session.begin()
                role = (
                    session.query(Role)
                    .options(joinedload(Role.permissions))
                    .filter(Role.uuid == role_uuid)
                    .first()
                )
                if not role:
                    raise NotFoundError(detail="Requested role does not exist")
                permissions_to_revoke = [
                    p for p in role.permissions
                    if p.uuid in permission_uuids
                ]
                if not permissions_to_revoke:
                    return role
                permission_uuids_to_revoke = [p.uuid for p in permissions_to_revoke]
                session.query(RolePermission).filter(
                    and_(
                        RolePermission.role_uuid == role_uuid,
                        RolePermission.permission_uuid.in_(permission_uuids_to_revoke),
                    )
                ).delete(synchronize_session=False)
                policies_to_remove = [
                    [role.uuid, permission.resource, permission.action, permission.module]
                    for permission in permissions_to_revoke
                ]
                with self._bulk_operation_context():
                    self.enforcer.remove_policies(policies_to_remove)
                session.commit()
                session.refresh(role)
                return role
            except Exception:
                raise

    def bulk_attach_permissions(self, role_uuid: str, permission_uuids: List[str]) -> Any:
        with self.db() as session:
            try:
                if not session.is_active:
                    session.begin()
                role = (
                    session.query(Role)
                    .options(joinedload(Role.permissions))
                    .filter(Role.uuid == role_uuid)
                    .first()
                )
                if not role:
                    raise NotFoundError(detail="Requested role does not exist")
                existing_permission_uuids = {p.uuid for p in role.permissions}
                new_permission_uuids = set(permission_uuids) - existing_permission_uuids
                if not new_permission_uuids:
                    return role
                new_permissions = (
                    session.query(Permission)
                    .filter(Permission.uuid.in_(new_permission_uuids))
                    .all()
                )
                if len(new_permissions) != len(new_permission_uuids):
                    found_permission_uuids = {p.uuid for p in new_permissions}
                    missing_permission_uuids = new_permission_uuids - found_permission_uuids
                    raise NotFoundError(
                        detail=f"Permissions with UUIDs '{', '.join(missing_permission_uuids)}' not found"
                    )
                role_permissions = [
                    {"role_uuid": role_uuid, "permission_uuid": p.uuid}
                    for p in new_permissions
                ]
                session.bulk_insert_mappings(RolePermission, role_permissions)
                policies_to_add = [
                    [role.uuid, permission.resource, permission.action, permission.module]
                    for permission in new_permissions
                ]
                with self._bulk_operation_context():
                    self.enforcer.add_policies(policies_to_add)
                session.commit()
                session.refresh(role)
                return role
            except Exception:
                raise

    def get_user_roles(self, user_uuid: str, session: Optional[Session] = None) -> List[Any]:
        def query_roles(session: Session) -> List[Any]:
            return (
                session.query(Role)
                .join(
                    UserRole,
                    and_(
                        UserRole.role_uuid == Role.uuid,
                        UserRole.user_uuid == user_uuid
                    )
                )
                .options(joinedload(Role.permissions))
                .all()
            )
        if session:
            return query_roles(session)
        else:
            with self.db() as new_session:
                return query_roles(new_session)

    def bulk_assign_roles_to_user(self, user_uuid: str, role_uuids: List[str]) -> List[Any]:
        with self.db() as session:
            try:
                if not session.is_active:
                    session.begin()
                current_roles = (
                    session.query(UserRole)
                    .options(joinedload(UserRole.role))
                    .filter(UserRole.user_uuid == user_uuid)
                    .all()
                )
                current_role_uuids = {role.role.uuid for role in current_roles}
                new_role_uuids = set(role_uuids) - current_role_uuids
                roles_to_remove = current_role_uuids - set(role_uuids)
                if roles_to_remove:
                    session.query(UserRole).filter(
                        and_(
                            UserRole.user_uuid == user_uuid,
                            UserRole.role_uuid.in_(roles_to_remove),
                        )
                    ).delete(synchronize_session=False)
                if new_role_uuids:
                    new_roles = (
                        session.query(Role).filter(Role.uuid.in_(new_role_uuids)).all()
                    )
                    if len(new_roles) != len(new_role_uuids):
                        raise NotFoundError(detail="One or more roles not found")
                    user_roles = [
                        UserRole(user_uuid=user_uuid, role_uuid=role.uuid)
                        for role in new_roles
                    ]
                    session.bulk_save_objects(user_roles)
                session.commit()
                return self.get_user_roles(user_uuid, session)
            except Exception:
                raise

    def bulk_revoke_roles_from_user(self, user_uuid: str, role_uuids: List[str]) -> List[Any]:
        with self.db() as session:
            try:
                if not session.is_active:
                    session.begin()
                current_roles = (
                    session.query(UserRole)
                    .options(joinedload(UserRole.role))
                    .filter(UserRole.user_uuid == user_uuid)
                    .filter(UserRole.role_uuid.in_(role_uuids))
                    .all()
                )
                if not current_roles:
                    return self.get_user_roles(user_uuid)
                role_uuids_to_revoke = {role.role.uuid for role in current_roles}
                session.query(UserRole).filter(
                    and_(
                        UserRole.user_uuid == user_uuid,
                        UserRole.role_uuid.in_(role_uuids_to_revoke),
                    )
                ).delete(synchronize_session=False)
                session.commit()
                return self.get_user_roles(user_uuid, session)
            except Exception:
                raise

    def bulk_attach_roles_to_user(self, user_uuid: str, role_uuids: List[str]) -> List[Any]:
        with self.db() as session:
            try:
                if not session.is_active:
                    session.begin()
                current_roles = (
                    session.query(UserRole)
                    .options(joinedload(UserRole.role))
                    .filter(UserRole.user_uuid == user_uuid)
                    .all()
                )
                current_role_uuids = {role.role.uuid for role in current_roles}
                new_role_uuids = set(role_uuids) - current_role_uuids
                if not new_role_uuids:
                    return self.get_user_roles(user_uuid)
                new_roles = (
                    session.query(Role).filter(Role.uuid.in_(new_role_uuids)).all()
                )
                if len(new_roles) != len(new_role_uuids):
                    raise NotFoundError(detail="There are some roles that does not exist.")
                user_roles = [
                    UserRole(user_uuid=user_uuid, role_uuid=role.uuid) for role in new_roles
                ]
                session.bulk_save_objects(user_roles)
                session.commit()
                return self.get_user_roles(user_uuid, session)
            except Exception:
                raise

    def check_permission(self, user_uuid: str, resource: str, action: str, module: str) -> bool:
        # ← NO CHANGES
        with self.db() as session:
            roles = (
                session.query(Role)
                .join(
                    UserRole,
                    and_(
                        UserRole.role_uuid == Role.uuid,
                        UserRole.user_uuid == user_uuid,
                    ),
                )
                .all()
            )
            for role in roles:
                if self.enforcer.enforce(role.uuid, resource, action, module):
                    return True
                if self.enforcer.enforce(role.name, resource, action, module):
                    return True
            return False

    def check_permission_by_role(self, role_uuid: str, resource: str, action: str, module: str) -> bool:
        # ← NO CHANGES
        if self.enforcer.enforce(role_uuid, resource, action, module):
            return True
        return False

    def check_permission_by_user(self, user_uuid: str, resource: str, action: str, module: str) -> bool:
        # ← NO CHANGES
        if self.enforcer.enforce(f"user:{user_uuid}", resource, action, module):
            return True
        return False

    def get_role(self, role_uuid: str, session: Optional[Session] = None) -> Any:
        def query_role(session: Session) -> Any:
            role = session.query(Role).filter(Role.uuid == role_uuid).first()
            if not role:
                raise NotFoundError(detail="Requested role does not exist.")
            return role
        if session:
            return query_role(session)
        else:
            with self.db() as session:
                return query_role(session)

    def is_watcher_active(self) -> bool:
        # ← CHANGED: Updated for Service Bus watcher
        if self.watcher is None:
            return False
        if isinstance(self.watcher, AzureServiceBusWatcher):
            return self.watcher._sender is not None
        # Fallback for any other watcher type
        return hasattr(self.watcher, 'pub_client') and self.watcher.pub_client is not None

    def get_watcher_status(self) -> dict:
        # ← CHANGED: Updated for Service Bus watcher
        if not self.is_watcher_active():
            return {
                "active": False,
                "message": "Watcher not initialized or not active"
            }

        if isinstance(self.watcher, AzureServiceBusWatcher):
            return {
                "active": True,
                "type": "azure_service_bus",
                "topic": self.watcher.config.topic_name,
                "subscription": self.watcher.subscription_name,
                "instance_id": self.watcher.instance_id,
                "connected": True,
            }

        return {"active": True, "type": "unknown"}

    def reload_policies(self) -> bool:
        try:
            self.enforcer.load_policy()
            return True
        except Exception as e:
            logger.error(f"Failed to reload policies: {e}")
            return False

    def get_policy_count(self) -> int:
        return len(self.enforcer.get_policy())

    def clear_all_policies(self) -> bool:
        try:
            self.enforcer.clear_policy()
            self.enforcer.save_policy()
            return True
        except Exception as e:
            logger.error(f"Failed to clear policies: {e}")
            return False