from rest_framework import mixins, permissions, status, viewsets from rest_framework.authtoken.models import Token from rest_framework.permissions import AllowAny, IsAuthenticated from rest_framework.response import Response from rest_framework.views import APIView from django.db.models import Count, Q from apps.library.models import MediaItem from .models import User from .serializers import ( E621CredentialsSerializer, RegisterSerializer, UserListSerializer, UserSerializer, UserUpdateSerializer, ) class RegisterView(APIView): permission_classes = [AllowAny] def post(self, request): serializer = RegisterSerializer(data=request.data) serializer.is_valid(raise_exception=True) user = serializer.save() token, _ = Token.objects.get_or_create(user=user) return Response( { "token": token.key, "user": UserSerializer(user, context={"request": request}).data, }, status=status.HTTP_201_CREATED, ) class MeView(APIView): permission_classes = [IsAuthenticated] def get(self, request): return Response(UserSerializer(request.user, context={"request": request}).data) class LogoutView(APIView): permission_classes = [IsAuthenticated] def post(self, request): if request.auth: request.auth.delete() return Response(status=status.HTTP_204_NO_CONTENT) class E621CredentialsView(APIView): """Store the e621 credentials used by the user's browser to call e621 directly.""" permission_classes = [IsAuthenticated] def _payload(self, user): return { "username": user.e621_username, "api_key": user.e621_api_key, "base_url": user.e621_base_url, "configured": user.e621_configured, } def get(self, request): return Response(self._payload(request.user)) def put(self, request): serializer = E621CredentialsSerializer(data=request.data) serializer.is_valid(raise_exception=True) data = serializer.validated_data user = request.user user.e621_username = data["username"].strip() user.e621_api_key = data["api_key"].strip() user.e621_base_url = ( (data.get("base_url") or "https://e621.net").strip().rstrip("/") ) user.save( update_fields=["e621_username", "e621_api_key", "e621_base_url"] ) return Response(self._payload(user)) class IsStaffUser(permissions.BasePermission): message = "Staff only." def has_permission(self, request, view): user = request.user return bool( user and user.is_authenticated and ( user.is_superuser or user.is_staff or user.role == user.ROLE_STAFF ) ) class UserViewSet( mixins.ListModelMixin, mixins.RetrieveModelMixin, viewsets.GenericViewSet, ): """Staff user directory: roles and J-ID avatars.""" permission_classes = [IsStaffUser] http_method_names = ["get", "patch", "head", "options"] def get_queryset(self): queryset = User.objects.annotate(uploads_count=Count("uploads")).order_by("id") search = self.request.query_params.get("search", "").strip() if search: queryset = queryset.filter( Q(username__icontains=search) | Q(email__icontains=search) ) return queryset def get_serializer_class(self): if self.request.method in {"PATCH", "PUT"}: return UserUpdateSerializer return UserListSerializer def update(self, request, *args, **kwargs): user = self.get_object() serializer = UserUpdateSerializer(data=request.data) serializer.is_valid(raise_exception=True) data = serializer.validated_data update_fields = [] if "role" in data: user.role = data["role"] update_fields.append("role") if "avatar_j_id" in data: value = str(data.get("avatar_j_id") or "").strip() if not value: user.avatar = None else: numeric = value[2:] if value.upper().startswith("J-") else value item = ( MediaItem.objects.filter(pk=int(numeric)).first() if numeric.isdigit() else None ) if item is None: return Response( {"detail": f"No library item {value}."}, status=status.HTTP_400_BAD_REQUEST, ) user.avatar = item update_fields.append("avatar") if update_fields: user.save(update_fields=update_fields) return Response( UserListSerializer(user, context=self.get_serializer_context()).data ) partial_update = update