trustme/core/apps/companies/views/companies.py

import uuid

from django_core.mixins import BaseViewSetMixin  # type: ignore
from drf_spectacular.utils import extend_schema
from rest_framework.decorators import action  # type: ignore
from rest_framework.permissions import AllowAny, IsAdminUser  # type: ignore
from rest_framework.viewsets import (  # type: ignore
    ModelViewSet, 
    GenericViewSet
)
from rest_framework.request import HttpRequest  # type: ignore
from rest_framework.response import Response  # type: ignore
from rest_framework import status  # type: ignore
from rest_framework.generics import get_object_or_404  # type: ignore


from django.contrib.auth import get_user_model
from core.apps.companies.permissions import IsCompanyAccount

from core.apps.companies.models import (
    CompanyModel,
    CompanyFolderModel,
    CompanyAccountModel
)
from core.apps.companies.serializers import (
    CreateCompanySerializer,
    ListCompanySerializer,
    RetrieveCompanySerializer,
    UpdateCompanySerializer,
    DestroyCompanySerializer,

    RetrieveCompanyFolderSerializer,
    RetrieveCompanyAccountSerializer,
    CreateCompanyFolderSerializer,
    BaseCompanyAccountSerializer,

    CreateFolderForCompanySerializer
)

from core.apps.contracts.serializers import (
    RetrieveContractSerializer,
    BaseContractSerializer
)

from core.apps.contracts.models import (
    ContractModel,
)


UserModel = get_user_model()


######################################################################
# View Only For Company CRUD That Belongs To Admin.
######################################################################
@extend_schema(tags=["Company"])
class CompanyView(BaseViewSetMixin, ModelViewSet):
    queryset = CompanyModel.objects.all()
    serializer_class = ListCompanySerializer
    permission_classes = [AllowAny]

    action_permission_classes = {
        "list": [IsAdminUser],
        "retrieve": [IsAdminUser],
        "create": [IsAdminUser],
        "update": [IsAdminUser],
        "destroy": [IsAdminUser],
    }
    action_serializer_class = {  # type: ignore
        "list": ListCompanySerializer,
        "retrieve": RetrieveCompanySerializer,
        "create": CreateCompanySerializer,
        "update": UpdateCompanySerializer,
        "destroy": DestroyCompanySerializer,
    }


######################################################################
# company/<uuid:pk>/contract Views
######################################################################
class CompanyContractViewSet(BaseViewSetMixin, GenericViewSet):
    queryset = CompanyModel.objects.all()
    permission_classes = [AllowAny]
    serializer_class = BaseContractSerializer

    action_permission_classes = {
        "list_contract": [IsCompanyAccount]
    }
    action_serializer_class = {
        "list_contract": RetrieveContractSerializer
    }

    #! TODO: status should be added.
    @extend_schema(
        summary="Company Contracts",
        description="Get List Company Contracts"
    )
    @action(methods=["GET"], detail=True, url_path="contracts")
    def list_contract(
        self,
        request: HttpRequest,
        *args: object, 
        **kwargs: object,
    ) -> Response:
        company = self.get_object()
        contracts = ContractModel.objects.filter(
            owners__legal_entity__phone=company.phone,
        ).distinct()

        folders_param = request.data.get("folders")
        if folders_param:
            folder_ids = folders_param.split(",")
            contracts = contracts.filter(folders__id__in=folder_ids)

        serializer = self.get_serializer(instance=contracts, many=True)  # type: ignore
        return Response(serializer.data, status=status.HTTP_200_OK)


######################################################################
# company/<uuid:pk>/accounts Views
######################################################################
class CompanyAccountViewSet(BaseViewSetMixin, GenericViewSet):
    queryset = CompanyModel.objects.all()
    serializer_class = BaseCompanyAccountSerializer
    permission_classes = [AllowAny]

    action_permission_classes = {
        "list_account": [IsCompanyAccount]
    }
    action_serializer_class = {
        "list_account": RetrieveCompanyAccountSerializer
    }

    @extend_schema(
        summary="List company accounts",
        description="List Company Accounts"
    )
    @action(url_path="accounts", detail=True, methods=["GET"])
    def list_account(
        self,
        request: HttpRequest,
        *args: object, 
        **kwargs: object,
    ) -> Response:
        company = self.get_object()
        accounts = CompanyAccountModel.objects.filter(company=company)
        ser = self.get_serializer(instance=accounts, many=True)  # type: ignore
        return Response(data=ser.data, status=status.HTTP_200_OK)


######################################################################
# company/<uuid:pk>/folders Views
######################################################################
class CompanyFolderViewSet(BaseViewSetMixin, GenericViewSet):
    queryset = CompanyModel.objects.all()
    permission_classes = [AllowAny]

    action_permission_classes = {
        "list_folder": [IsCompanyAccount],
        "create_folder": [IsCompanyAccount],
    }
    action_serializer_class = {  # type: ignore
        "list_folder": RetrieveCompanyFolderSerializer, 
        "create_folder": CreateFolderForCompanySerializer,
    }

    @extend_schema(
        summary="List Company Folders",
        description="List Company Folders"
    )
    @action(methods=["GET"], detail=True, url_path="folders")
    def list_folder(self, *args: object, **kwargs: object) -> Response:
        company = self.get_object()
        folders = CompanyFolderModel.objects.filter(company=company)
        ser = self.get_serializer(
            instance=folders, many=True, context={"company_id": company.pk}
        )  # type: ignore
        return Response(data=ser.data, status=status.HTTP_200_OK)
    
    @extend_schema(
        summary="Create Folder for company",
        description="Create Folder for company",
    )
    @action(url_path="folders", detail=True, methods=["POST"])
    def create_folder(self, request: HttpRequest, *args: object, **kwargs: object) -> Response:
        company = self.get_object()
        data = request.data.copy() | dict(company=company.id)  # type: ignore
        ser = self.get_serializer(data=data)  # type: ignore
        ser.is_valid(raise_exception=True)
        return Response(data=ser.data, status=status.HTTP_201_CREATED)