J621/backend/apps/library/uploads.py

"""Staged uploads: files land in a temp folder until they are resolved.

Flow:
- POST /api/uploads/            stage a file (computes MD5, detects duplicates)
- GET  /api/uploads/            list my staged uploads
- GET  /api/uploads/<id>/file/  serve the staged file (previews, IQDB)
- POST /api/uploads/<id>/iqdb/  store IQDB candidates -> visual_match
- POST /api/uploads/<id>/resolve/ link to a post or apply custom metadata
- DELETE /api/uploads/<id>/     discard a staged upload / dismiss a record
"""

import logging
import shutil
from pathlib import Path

from django.conf import settings
from django.http import Http404
from rest_framework import mixins, status, viewsets
from rest_framework.decorators import action
from rest_framework.parsers import FormParser, JSONParser, MultiPartParser
from rest_framework.response import Response

from . import services
from .models import MediaItem, TempUpload
from .permissions import CanUpload
from .serializers import TempUploadSerializer

logger = logging.getLogger(__name__)


def complete_temp_upload(temp):
    """Move a staged file into the watched folder and index it."""
    folder = Path(settings.WATCHED_FOLDER)
    if not temp.file:
        raise ValueError("staged file is missing")

    destination = services.unique_destination(
        folder, temp.original_filename or "upload"
    )
    with temp.file.open("rb") as source, open(destination, "wb") as target:
        shutil.copyfileobj(source, target)

    item, _, _ = services.index_file(destination, folder)
    temp.file.delete(save=False)
    temp.library_item = item
    temp.status = TempUpload.STATUS_COMPLETED

    update_fields = []
    if item.uploaded_by_id is None:
        item.uploaded_by = temp.user
        update_fields.append("uploaded_by")
    if temp.e621_post_id is not None:
        item.e621_post_id = temp.e621_post_id
        item.e621_data = temp.e621_data
        update_fields += ["e621_post_id", "e621_data"]
    if temp.custom_rating in {"s", "q", "e"} and not item.rating:
        item.rating = temp.custom_rating
        update_fields.append("rating")
    if temp.custom_tags:
        item.tags = temp.custom_tags
        update_fields.append("tags")
    if temp.custom_notes:
        item.notes = temp.custom_notes
        update_fields.append("notes")
    if update_fields:
        item.save(update_fields=update_fields + ["updated_at"])

    temp.save()
    return item


class TempUploadViewSet(
    mixins.ListModelMixin,
    mixins.RetrieveModelMixin,
    mixins.DestroyModelMixin,
    viewsets.GenericViewSet,
):
    serializer_class = TempUploadSerializer
    permission_classes = [CanUpload]
    parser_classes = [MultiPartParser, FormParser, JSONParser]
    http_method_names = ["get", "post", "delete", "head", "options"]

    def get_queryset(self):
        queryset = TempUpload.objects.select_related("library_item")
        user = self.request.user
        if not (user.is_staff or user.is_superuser):
            queryset = queryset.filter(user=user)
        return queryset

    def create(self, request):
        upload = request.FILES.get("file")
        if upload is None:
            return Response(
                {"detail": "A file is required."}, status=status.HTTP_400_BAD_REQUEST
            )
        extension = Path(upload.name).suffix.lower()
        if extension not in services.ALLOWED_EXTENSIONS:
            return Response(
                {"detail": f"Unsupported file type: {extension or 'unknown'}"},
                status=status.HTTP_400_BAD_REQUEST,
            )

        temp = TempUpload.objects.create(
            user=request.user,
            file=upload,
            original_filename=upload.name,
            size=upload.size,
        )
        temp.md5 = services.compute_md5(temp.file.path)

        existing = MediaItem.objects.filter(md5=temp.md5).first()
        if existing is not None:
            temp.status = TempUpload.STATUS_COMPLETED
            temp.resolution = TempUpload.RESOLUTION_DUPLICATE
            temp.library_item = existing
            temp.file.delete(save=False)
        temp.save()
        return Response(
            TempUploadSerializer(temp).data, status=status.HTTP_201_CREATED
        )

    @action(detail=True, methods=["get"])
    def file(self, request, pk=None):
        temp = self.get_object()
        if not temp.file:
            raise Http404
        return services.serve_file(request, temp.file.path)

    @action(detail=True, methods=["post"])
    def iqdb(self, request, pk=None):
        temp = self.get_object()
        results = request.data.get("results")
        if not isinstance(results, list):
            return Response(
                {"detail": "results must be a list."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        temp.iqdb_data = services.sanitize_iqdb_results(results)
        if temp.status == TempUpload.STATUS_PENDING and temp.iqdb_data:
            temp.status = TempUpload.STATUS_VISUAL_MATCH
        temp.save(update_fields=["iqdb_data", "status", "updated_at"])
        return Response(TempUploadSerializer(temp).data)

    @action(detail=True, methods=["post"])
    def resolve(self, request, pk=None):
        temp = self.get_object()
        if temp.status == TempUpload.STATUS_COMPLETED:
            return Response(
                {"detail": "This upload is already in the library."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        mode = str(request.data.get("mode") or "").strip()
        if mode == "link":
            post = request.data.get("post")
            post_id = request.data.get("post_id")
            trimmed = services.trim_e621_post(post)
            has_numeric_id = post_id is not None and str(post_id).isdigit()
            if trimmed is None and not has_numeric_id:
                return Response(
                    {"detail": "A post payload or post_id is required."},
                    status=status.HTTP_400_BAD_REQUEST,
                )
            if trimmed is not None and trimmed.get("id") is not None:
                temp.e621_post_id = int(trimmed["id"])
            elif has_numeric_id:
                temp.e621_post_id = int(post_id)
            temp.e621_data = trimmed
            temp.resolution = (
                TempUpload.RESOLUTION_AUTO_MD5
                if request.data.get("auto")
                else TempUpload.RESOLUTION_LINKED
            )
            rating = str(request.data.get("rating") or "")
            if rating in {"s", "q", "e"}:
                temp.custom_rating = rating
            tags = services.parse_tags(request.data.get("tags"))
            if tags is not None:
                temp.custom_tags = tags
            notes = request.data.get("notes")
            if notes is not None:
                temp.custom_notes = str(notes)
        elif mode == "custom":
            rating = str(request.data.get("rating") or "")
            temp.custom_rating = rating if rating in {"s", "q", "e"} else ""
            temp.custom_tags = services.parse_tags(request.data.get("tags")) or []
            temp.custom_notes = str(request.data.get("notes") or "")
            temp.resolution = TempUpload.RESOLUTION_CUSTOM
        else:
            return Response(
                {"detail": "mode must be 'link' or 'custom'."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        temp.save()
        try:
            complete_temp_upload(temp)
        except Exception as exc:  # noqa: BLE001 - report completion failures
            logger.exception("Could not complete staged upload %s", temp.id)
            temp.status = TempUpload.STATUS_ERROR
            temp.save(update_fields=["status", "updated_at"])
            return Response(
                {"detail": f"Could not finish the upload: {exc}"},
                status=status.HTTP_400_BAD_REQUEST,
            )
        temp.refresh_from_db()
        return Response(TempUploadSerializer(temp).data)

    def perform_destroy(self, instance):
        if instance.file:
            instance.file.delete(save=False)
        instance.delete()