J621/backend/apps/accounts/views.py

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