Dynavera/apps/domains/viewsets.py

from rest_framework.viewsets import ModelViewSet
from rest_framework.permissions import IsAuthenticated, IsAuthenticatedOrReadOnly
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework import status
from django.shortcuts import get_object_or_404
from django.utils import timezone
from apps.domains.models import Domain, Organization, Dataset, OrganizationMembership, InviteToken, DomainMembership
from apps.domains.serializers import (
    DomainSerializer,
    OrganizationSerializer,
    DatasetSerializer,
    OrganizationMembershipSerializer,
    InviteTokenSerializer,
    DomainMembershipSerializer,
)


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

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

    def perform_create(self, serializer):
        org = serializer.save(owner=self.request.user)
        OrganizationMembership.objects.create(
            organization=org,
            user=self.request.user,
            role=OrganizationMembership.Role.EMPLOYER
        )

    def update(self, request, *args, **kwargs):
        org = self.get_object()
        membership = OrganizationMembership.objects.filter(
            organization=org,
            user=request.user,
            role=OrganizationMembership.Role.EMPLOYER
        ).first()
        if not membership:
            return Response(
                {"error": "Only employers can update organization details"},
                status=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,
            role=OrganizationMembership.Role.EMPLOYER
        ).first()
        if not membership:
            return Response(
                {"error": "Only employers can update member roles"},
                status=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=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,
            role=OrganizationMembership.Role.EMPLOYER
        ).first()
        if not membership:
            return Response(
                {"error": "Only employers can remove members"},
                status=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=status.HTTP_400_BAD_REQUEST
            )
        target_membership.delete()
        return Response(status=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 = InviteTokenSerializer(tokens, many=True, context={'request': request})
            return Response(serializer.data)

        elif request.method == 'POST':
            membership = OrganizationMembership.objects.filter(
                organization=org,
                user=request.user,
                role=OrganizationMembership.Role.EMPLOYER
            ).first()
            if not membership:
                return Response(
                    {"error": "Only employers can create invites"},
                    status=status.HTTP_403_FORBIDDEN
                )

            token = InviteToken.objects.create(
                organization=org,
                created_by=request.user
            )
            serializer = InviteTokenSerializer(token, context={'request': request})
            return Response(serializer.data, status=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,
            role=OrganizationMembership.Role.EMPLOYER
        ).first()
        if not membership:
            return Response(
                {"error": "Only employers can revoke invites"},
                status=status.HTTP_403_FORBIDDEN
            )

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

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

    @action(detail=True, methods=['get'], url_path='domains/(?P<domain_id>[^/.]+)/members')
    def domain_members(self, request, uuid=None, domain_id=None):
        org = self.get_object()
        domain = get_object_or_404(Domain, organization=org, id=domain_id)
        memberships = domain.memberships.all()
        serializer = DomainMembershipSerializer(memberships, many=True)
        return Response(serializer.data)

    @action(detail=True, methods=['post'], url_path='domains/(?P<domain_id>[^/.]+)/members')
    def add_domain_member(self, request, uuid=None, domain_id=None):
        org = self.get_object()
        domain = get_object_or_404(Domain, organization=org, id=domain_id)

        user_id = request.data.get('user_id')
        org_membership = OrganizationMembership.objects.filter(
            organization=org,
            user_id=user_id
        ).first()

        if not org_membership:
            return Response(
                {"error": "User must be a member of the organization first"},
                status=status.HTTP_400_BAD_REQUEST
            )

        domain_membership, created = DomainMembership.objects.get_or_create(
            domain=domain,
            user_id=user_id
        )

        serializer = DomainMembershipSerializer(domain_membership)
        return Response(serializer.data, status=status.HTTP_201_CREATED if created else status.HTTP_200_OK)


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

    def get_queryset(self):
        return InviteToken.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=status.HTTP_400_BAD_REQUEST
            )

        membership, created = OrganizationMembership.objects.get_or_create(
            organization=invite.organization,
            user=request.user,
            defaults={'role': OrganizationMembership.Role.EMPLOYEE}
        )

        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=status.HTTP_201_CREATED if created else status.HTTP_200_OK)


class DomainViewSet(ModelViewSet):
    queryset = Domain.objects.all()
    serializer_class = DomainSerializer
    permission_classes = [IsAuthenticatedOrReadOnly]
    lookup_field = 'uuid'

    def get_queryset(self):
        user = self.request.user
        if user.is_authenticated:
            return Domain.objects.filter(
                organization__memberships__user=user
            ).distinct()
        return Domain.objects.none()


class DatasetViewSet(ModelViewSet):
    queryset = Dataset.objects.all()
    serializer_class = DatasetSerializer
    permission_classes = [IsAuthenticatedOrReadOnly]
    lookup_field = 'uuid'