Dynavera/apps/orgs/viewsets.py

from rest_framework.viewsets import ModelViewSet
from rest_framework.permissions import IsAuthenticated
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.status import HTTP_200_OK, HTTP_201_CREATED, HTTP_204_NO_CONTENT, HTTP_400_BAD_REQUEST, HTTP_403_FORBIDDEN
from django.shortcuts import get_object_or_404
from django.utils import timezone
from django.db.models import Q
from apps.orgs.models import Organization, OrganizationMembership, OrganizationInvitation, Role, RoleMembership
from apps.orgs.serializers import (
    OrganizationSerializer,
    OrganizationMembershipSerializer,
    OrganizationInvitationSerializer,
    RoleSerializer,
    RoleMembershipSerializer,
)

class OrganizationViewSet(ModelViewSet):
    queryset = Organization.objects.all()
    serializer_class = OrganizationSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'uuid'

    def get_queryset(self):
        return Organization.objects.filter(Q(memberships__user=self.request.user) | Q(owner=self.request.user)).distinct()

    def perform_create(self, serializer):
        org = serializer.save(owner = self.request.user)
        OrganizationMembership.objects.create(organization = org, user = self.request.user, is_manager = True)

    def update(self, request, *args, **kwargs):
        org = self.get_object()
        membership = OrganizationMembership.objects.filter(
            organization=org, user=request.user, is_manager=True
        ).first()
        if not membership:
            return Response({'error': 'Only managers can update organization details'}, status=HTTP_403_FORBIDDEN)
        return super().update(request, *args, **kwargs)

    @action(detail=True, methods=['get'])
    def members(self, request, uuid=None):
        org = self.get_object()
        memberships = org.memberships.all()
        serializer = OrganizationMembershipSerializer(memberships, many=True)
        return Response(serializer.data)

    @action(detail=True, methods=['patch'], url_path='members/(?P<user_id>[^/.]+)')
    def update_member(self, request, uuid=None, user_id=None):
        org = self.get_object()
        membership = OrganizationMembership.objects.filter(
            organization=org, user=request.user, is_manager=True
        ).first()
        if not membership:
            return Response({'error': 'Only managers can update member roles'}, status=HTTP_403_FORBIDDEN)

        target_membership = get_object_or_404(OrganizationMembership, organization=org, user_id=user_id)
        serializer = OrganizationMembershipSerializer(target_membership, data=request.data, partial = True)
        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data)
        return Response(serializer.errors, status=HTTP_400_BAD_REQUEST)

    @action(detail=True, methods=['delete'], url_path='members/(?P<user_id>[^/.]+)')
    def remove_member(self, request, uuid=None, user_id=None):
        org = self.get_object()
        membership = OrganizationMembership.objects.filter(
            organization=org, user=request.user, is_manager=True
        ).first()
        if not membership:
            return Response({'error': 'Only managers can remove members'}, status=HTTP_403_FORBIDDEN)

        target_membership = get_object_or_404(OrganizationMembership, organization=org, user_id=user_id)
        if target_membership.user == org.owner:
            return Response({'error': 'Cannot remove the organization owner'}, status=HTTP_400_BAD_REQUEST)
        target_membership.delete()
        return Response(status=HTTP_204_NO_CONTENT)

    @action(detail=True, methods=['get', 'post'])
    def invites(self, request, uuid=None):
        org = self.get_object()

        if request.method == 'GET':
            tokens = org.invite_tokens.filter(is_active = True, used_by__isnull = True)
            serializer = OrganizationInvitationSerializer(tokens, many = True, context = {'request': request})
            return Response(serializer.data)

        membership = OrganizationMembership.objects.filter(organization=org, user=request.user, is_manager=True).first()
        if not membership:
            return Response({'error': 'Only managers can create invites'}, status=HTTP_403_FORBIDDEN)

        token = OrganizationInvitation.objects.create(organization = org, created_by = request.user)
        serializer = OrganizationInvitationSerializer(token, context = {'request': request})
        return Response(serializer.data, status = HTTP_201_CREATED)

    @action(detail=True, methods=['delete'], url_path='invites/(?P<token>[^/.]+)')
    def revoke_invite(self, request, uuid=None, token=None):
        org = self.get_object()
        membership = OrganizationMembership.objects.filter(organization=org, user=request.user, is_manager=True).first()
        if not membership:
            return Response({'error': 'Only managers can revoke invites'}, status=HTTP_403_FORBIDDEN)

        invite = get_object_or_404(OrganizationInvitation, organization=org, token=token)
        invite.is_active = False
        invite.save()
        return Response(status=HTTP_204_NO_CONTENT)

    @action(detail=True, methods=['get', 'post'], url_path='role')
    def role(self, request, uuid=None):
        org = self.get_object()

        if request.method == 'GET':
            roles = Role.objects.filter(organization=org)
            serializer = RoleSerializer(roles, many=True)
            return Response(serializer.data)

        membership = OrganizationMembership.objects.filter(organization=org, user=request.user, is_manager=True).first()
        if not membership:
            return Response({'error': 'Only managers can create roles'}, status=HTTP_403_FORBIDDEN)

        serializer = RoleSerializer(data=request.data)
        if serializer.is_valid():
            serializer.save(organization=org)
            return Response(serializer.data, status=HTTP_201_CREATED)
        return Response(serializer.errors, status=HTTP_400_BAD_REQUEST)

    @action(detail=True, methods=['get', 'post'], url_path='role/(?P<role_id>[^/.]+)/members')
    def role_members(self, request, uuid=None, role_id=None):
        org = self.get_object()
        role = get_object_or_404(Role, id=role_id, organization=org)
        requester_membership = OrganizationMembership.objects.filter(organization=org, user=request.user).first()
        if not requester_membership:
            return Response(status=HTTP_404_NOT_FOUND)

        if request.method == 'GET':
            memberships = RoleMembership.objects.filter(role=role)
            serializer = RoleMembershipSerializer(memberships, many=True)
            return Response(serializer.data)
        manager_membership = OrganizationMembership.objects.filter(organization=org, user=request.user, is_manager=True).first()
        user_id = request.data.get('user_id')
        if not user_id:
            return Response({'error': 'user_id is required'}, status=HTTP_400_BAD_REQUEST)
        if request.user.id != int(user_id) and not manager_membership:
            return Response({'error': 'Only managers can add other users to roles'}, status=HTTP_403_FORBIDDEN)

        role_membership, created = RoleMembership.objects.get_or_create(role=role, user_id=user_id)

        serializer = RoleMembershipSerializer(role_membership)
        return Response(serializer.data, status=HTTP_201_CREATED if created else HTTP_200_OK)

class InviteViewSet(ModelViewSet):
    queryset = OrganizationInvitation.objects.all()
    serializer_class = OrganizationInvitationSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'token'
    http_method_names = ['get', 'post', 'delete']

    def get_queryset(self):
        return OrganizationInvitation.objects.filter(is_active = True, used_by__isnull = True)

    @action(detail=True, methods=['post'])
    def accept(self, request, token=None):
        invite = self.get_object()

        if not invite.is_valid():
            return Response({'error': 'This invite is no longer valid'}, status = HTTP_400_BAD_REQUEST)
        membership, created = OrganizationMembership.objects.get_or_create(organization = invite.organization, user = request.user, defaults = {'is_manager': False})
        if created:
            invite.used_by = request.user
            invite.used_at = timezone.now()
            invite.is_active = False
            invite.save()
        serializer = OrganizationSerializer(invite.organization)
        return Response(serializer.data, status = HTTP_201_CREATED if created else HTTP_200_OK)

class RoleViewSet(ModelViewSet):
    queryset = Role.objects.all()
    serializer_class = RoleSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'uuid'

    def get_queryset(self):
        return Role.objects.filter(Q(organization__memberships__user=self.request.user) | Q(organization__owner=self.request.user)).distinct()

    def perform_create(self, serializer):
        serializer.save()