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
            )
        )


def resolve_avatar_item(value):
    """Turn a "J-42" / "42" string into a MediaItem.

    Returns ``(item, error)``: an empty value clears the avatar (``None``,
    ``None``), an unknown item returns an error message.
    """
    value = str(value or "").strip()
    if not value:
        return None, None
    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 None, f"No library item {value}."
    return item, None


class AvatarView(APIView):
    """Self-service profile picture chosen from the library."""

    permission_classes = [IsAuthenticated]

    def post(self, request):
        serializer = UserUpdateSerializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        item, error = resolve_avatar_item(
            serializer.validated_data.get("avatar_j_id")
        )
        if error:
            return Response({"detail": error}, status=status.HTTP_400_BAD_REQUEST)
        user = request.user
        user.avatar = item
        user.save(update_fields=["avatar"])
        return Response(UserSerializer(user, context={"request": request}).data)


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:
            item, error = resolve_avatar_item(data.get("avatar_j_id"))
            if error:
                return Response(
                    {"detail": error},
                    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