arrr-erp/app/modules/core/iam/scope.py

from __future__ import annotations

from dataclasses import dataclass

from fastapi import HTTPException
from sqlalchemy import select
from sqlalchemy.orm import Session

from app.modules.core.iam.models import User
from app.modules.core.rbac.models import Role, UserRole
from app.modules.core.tenancy.models import Branch, Tenant

ROLE_SYSTEM_ADMIN = "System Admin"
ROLE_FIRM_ADMIN = "Firm Admin"
ROLE_PARTNER = "Partner"
ROLE_BRANCH_MANAGER = "Branch Manager"
ROLE_STAFF = "Staff"
ROLE_CLIENT = "Client"
ROLE_CONSULTANT = "Consultant"

# Matrix: only System Admin and Firm Admin manage users.
MANAGEABLE_ROLES_BY_ACTOR = {
    ROLE_SYSTEM_ADMIN: {
        ROLE_SYSTEM_ADMIN,
        ROLE_FIRM_ADMIN,
        ROLE_PARTNER,
        ROLE_BRANCH_MANAGER,
        ROLE_STAFF,
        ROLE_CLIENT,
        ROLE_CONSULTANT,
    },
    ROLE_FIRM_ADMIN: {
        ROLE_PARTNER,
        ROLE_BRANCH_MANAGER,
        ROLE_STAFF,
        ROLE_CLIENT,
        ROLE_CONSULTANT,
    },
}


@dataclass
class UserScope:
    actor: User
    role_names: list[str]
    is_system_admin: bool
    is_firm_admin: bool
    is_partner: bool
    is_branch_manager: bool

    @property
    def tenant_scoped(self) -> bool:
        return self.is_firm_admin or self.is_partner or self.is_branch_manager

    @property
    def branch_scoped(self) -> bool:
        return self.is_branch_manager


class ScopeError(Exception):
    pass


def get_role_names(db: Session, user_id: int) -> list[str]:
    q = (
        select(Role.name)
        .join(UserRole, UserRole.role_id == Role.id)
        .where(UserRole.user_id == user_id, Role.is_active.is_(True))
        .order_by(Role.name)
    )
    return [name for (name,) in db.execute(q).all()]


def build_scope(db: Session, actor: User) -> UserScope:
    role_names = get_role_names(db, actor.id)
    return UserScope(
        actor=actor,
        role_names=role_names,
        is_system_admin=ROLE_SYSTEM_ADMIN in role_names,
        is_firm_admin=ROLE_FIRM_ADMIN in role_names,
        is_partner=ROLE_PARTNER in role_names,
        is_branch_manager=ROLE_BRANCH_MANAGER in role_names,
    )


def ensure_users_view_scope(scope: UserScope) -> None:
    if scope.is_system_admin or scope.is_firm_admin or scope.is_partner or scope.is_branch_manager:
        return
    raise ScopeError("You are not allowed to view users.")


def ensure_users_manage_scope(scope: UserScope) -> None:
    if scope.is_system_admin or scope.is_firm_admin:
        return
    raise ScopeError("Only System Admin and Firm Admin can manage users.")


def list_visible_tenants(db: Session, scope: UserScope) -> list[Tenant]:
    if scope.is_system_admin:
        return db.execute(
            select(Tenant).where(Tenant.is_active.is_(True)).order_by(Tenant.name)
        ).scalars().all()
    tenant = db.execute(select(Tenant).where(Tenant.id == scope.actor.tenant_id)).scalar_one_or_none()
    return [tenant] if tenant else []


def list_visible_branches(db: Session, scope: UserScope, tenant_id: int | None = None) -> list[Branch]:
    effective_tenant_id = tenant_id or scope.actor.tenant_id
    q = select(Branch).where(Branch.is_active.is_(True), Branch.tenant_id == effective_tenant_id)
    if scope.branch_scoped:
        q = q.where(Branch.id == scope.actor.branch_id)
    return db.execute(q.order_by(Branch.name)).scalars().all()


def list_scoped_users(db: Session, scope: UserScope) -> list[User]:
    q = select(User)
    if not scope.is_system_admin:
        q = q.where(User.tenant_id == scope.actor.tenant_id)
    if scope.branch_scoped:
        q = q.where(User.branch_id == scope.actor.branch_id)
    return db.execute(q.order_by(User.id)).scalars().all()


def get_manageable_roles(db: Session, scope: UserScope) -> list[Role]:
    if scope.is_system_admin:
        return db.execute(
            select(Role).where(Role.is_active.is_(True)).order_by(Role.name)
        ).scalars().all()

    allowed_names: set[str] = set()
    for role_name in scope.role_names:
        allowed_names.update(MANAGEABLE_ROLES_BY_ACTOR.get(role_name, set()))
    if not allowed_names:
        return []
    return db.execute(
        select(Role).where(Role.is_active.is_(True), Role.name.in_(sorted(allowed_names))).order_by(Role.name)
    ).scalars().all()


def get_user_role_names(db: Session, user_id: int) -> list[str]:
    return get_role_names(db, user_id)


def get_user_role_ids(db: Session, user_id: int) -> list[int]:
    return db.execute(select(UserRole.role_id).where(UserRole.user_id == user_id)).scalars().all()


def can_manage_role_names(scope: UserScope, role_names: list[str]) -> bool:
    if scope.is_system_admin:
        return True
    allowed: set[str] = set()
    for actor_role in scope.role_names:
        allowed.update(MANAGEABLE_ROLES_BY_ACTOR.get(actor_role, set()))
    return set(role_names).issubset(allowed)


def validate_branch_matches_tenant(db: Session, tenant_id: int, branch_id: int) -> Branch:
    branch = db.execute(
        select(Branch).where(
            Branch.id == branch_id,
            Branch.tenant_id == tenant_id,
            Branch.is_active.is_(True),
        )
    ).scalar_one_or_none()
    if not branch:
        raise ScopeError("Selected branch does not belong to the selected tenant.")
    return branch


def resolve_target_tenant_branch(
    db: Session,
    scope: UserScope,
    tenant_id: int | None,
    branch_id: int | None,
) -> tuple[int, int]:
    if scope.is_system_admin:
        if tenant_id is None or branch_id is None:
            raise ScopeError("Tenant and branch are required.")
        validate_branch_matches_tenant(db, tenant_id, branch_id)
        return tenant_id, branch_id

    effective_tenant_id = scope.actor.tenant_id
    effective_branch_id = branch_id

    if tenant_id is not None and tenant_id != scope.actor.tenant_id:
        raise ScopeError("Cross-tenant user creation is not allowed.")

    if effective_branch_id is None:
        raise ScopeError("Branch is required.")

    validate_branch_matches_tenant(db, effective_tenant_id, effective_branch_id)
    return effective_tenant_id, effective_branch_id


def ensure_manageable_existing_user(db: Session, scope: UserScope, target_user: User) -> None:
    if scope.is_system_admin:
        return
    if not scope.is_firm_admin:
        raise ScopeError("Only System Admin and Firm Admin can manage users.")
    if target_user.tenant_id != scope.actor.tenant_id:
        raise ScopeError("You cannot manage users of another tenant.")

    target_roles = get_user_role_names(db, target_user.id)
    if target_roles and not can_manage_role_names(scope, target_roles):
        raise ScopeError("You cannot manage the selected user's role level.")


def ensure_assignable_roles(db: Session, scope: UserScope, role_ids: list[int]) -> list[Role]:
    if not role_ids:
        return []
    roles = db.execute(
        select(Role).where(Role.id.in_(role_ids), Role.is_active.is_(True)).order_by(Role.name)
    ).scalars().all()
    if len(roles) != len(set(role_ids)):
        raise ScopeError("One or more selected roles are invalid.")
    if not can_manage_role_names(scope, [r.name for r in roles]):
        raise ScopeError("You cannot assign one or more selected roles.")
    return roles


def assert_can_manage_role_object(scope: UserScope, role: Role) -> None:
    if scope.is_system_admin:
        return
    raise ScopeError("Only System Admin can manage RBAC roles.")


def scope_to_http(exc: ScopeError) -> HTTPException:
    return HTTPException(status_code=403, detail=str(exc))