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 urllib.parse import urlparse

from django.conf import settings
from django.contrib.auth import get_user_model
from django.core import signing
from django.core.exceptions import ValidationError
from django.http import Http404
from django.utils import timezone
from rest_framework import mixins, status, viewsets
from rest_framework.decorators import action
from rest_framework.parsers import FormParser, JSONParser, MultiPartParser
from rest_framework.permissions import AllowAny
from rest_framework.response import Response

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

logger = logging.getLogger(__name__)


def find_library_matches(path, limit=10, user=None, request=None):
    """Library items visually similar to a staged file."""
    hashes = services.compute_visual_hashes(path)
    if not hashes:
        return []
    algorithms = list(HASH_FIELDS)
    threshold = settings.VISUAL_MATCH_THRESHOLD
    matches = []
    for item in MediaItem.objects.prefetch_related("locations"):
        similarity = hashes_similarity(
            hashes,
            {field: getattr(item, field, "") for field in algorithms},
            algorithms,
            threshold,
        )
        if similarity is None:
            continue
        location = item.locations.first()
        matches.append(
            {
                "j_id": f"J-{item.id}",
                "filename": Path(location.rel_path).name if location else item.md5,
                "similarity": round(similarity * 100, 1),
                "thumbnail_url": services.signed_media_url(
                    item, user, "thumbnail", request=request
                ),
            }
        )
    matches.sort(key=lambda entry: entry["similarity"], reverse=True)
    return matches[:limit]


def complete_temp_upload(temp, download_url=None):
    """Index the upload into the library.

    With a ``download_url`` the e621 original is fetched and the staged copy
    is dropped; otherwise the staged file itself moves into the library.
    """
    folder = Path(settings.WATCHED_FOLDER)

    if download_url:
        file_data = {}
        if isinstance(temp.e621_data, dict):
            file_data = temp.e621_data.get("file") or {}
        extension = Path(urlparse(download_url).path).suffix
        if not extension:
            extension = Path(temp.original_filename).suffix
        md5 = file_data.get("md5")
        name = f"{md5}{extension}" if md5 else (temp.original_filename or "download")
        destination = services.unique_destination(folder, name)
        try:
            services.download_file(download_url, destination)
        except Exception:
            destination.unlink(missing_ok=True)
            raise
        item, _, location, _ = services.index_file(destination, folder)
        services.rename_location_to_j_id(item, location)
        services.ensure_visual_hashes(item)
        if temp.file:
            temp.file.delete(save=False)
    else:
        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, _, location, _ = services.index_file(destination, folder)
        services.rename_location_to_j_id(item, location)
        services.ensure_visual_hashes(item)
        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
        item.e621_match_status = MediaItem.E621_MATCHED
        item.e621_checked_at = timezone.now()
        update_fields += [
            "e621_post_id",
            "e621_data",
            "e621_match_status",
            "e621_checked_at",
        ]
    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]
    # Unpaginated: the board shows every staged upload (69-file batches were
    # silently cut to the API's 48-item page).
    pagination_class = None
    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_app_staff:
            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)
        # Visual similarity is deliberately a separate phase (the
        # /visual-match action) so a large batch uploads at full speed and
        # the board runs MD5 -> visual -> IQDB over the whole batch.
        temp.save()
        return Response(
            self.get_serializer(temp).data, status=status.HTTP_201_CREATED
        )

    @action(detail=True, methods=["post"], url_path="visual-match")
    def visual_match(self, request, pk=None):
        """Run the local visual-similarity pass for one staged upload."""
        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,
            )
        if not temp.file:
            return Response(
                {"detail": "The staged file is missing."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        matches = find_library_matches(
            temp.file.path, user=request.user, request=request
        )
        temp.visual_matches = matches
        if matches and temp.status == TempUpload.STATUS_PENDING:
            temp.status = TempUpload.STATUS_VISUAL_MATCH
        temp.save(update_fields=["visual_matches", "status", "updated_at"])
        return Response(self.get_serializer(temp).data)

    @action(detail=True, methods=["get", "head"], permission_classes=[AllowAny])
    def file(self, request, pk=None):
        """Serve the staged file; accepts a signed URL for media tags."""
        user = request.user if request.user.is_authenticated else None
        if user is None:
            signature = request.query_params.get("sig")
            if signature:
                try:
                    payload = signing.loads(
                        signature,
                        salt=services.UPLOAD_FILE_SALT,
                        max_age=86400,
                    )
                except signing.BadSignature:
                    payload = None
                if payload and str(payload.get("temp")) == str(pk):
                    user = (
                        get_user_model()
                        .objects.filter(pk=payload.get("user"))
                        .first()
                    )
        if user is None or not user.can_upload:
            return Response(
                {"detail": "Authentication required."},
                status=status.HTTP_401_UNAUTHORIZED,
            )

        temp = TempUpload.objects.filter(pk=pk).first()
        is_owner = temp is not None and temp.user_id == user.id
        if temp is None or not (is_owner or user.is_app_staff):
            raise Http404
        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(self.get_serializer(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()
        download_url = None
        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)

            # Prefer the e621 original; fall back to the staged file when the
            # post has no URL or its file is byte-identical to the upload.
            file_data = (trimmed or {}).get("file") or {}
            candidate_url = str(
                request.data.get("file_url") or file_data.get("url") or ""
            ).strip()
            post_md5 = str(file_data.get("md5") or "").strip().lower()
            if candidate_url and post_md5 and post_md5 == temp.md5.lower():
                candidate_url = ""
            download_url = candidate_url or None
        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, download_url=download_url)
        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(self.get_serializer(temp).data)

    @action(detail=False, methods=["post"], url_path="link-bulk")
    def link_bulk(self, request):
        """Attach e621 posts to many staged uploads in one request.

        The upload board's MD5 phase sends one posts.json query per 75 files
        and hands the matches over here, so a whole batch moves into the
        library in one update instead of one request per file.
        """
        links = request.data.get("links")
        if not isinstance(links, list) or not links:
            return Response(
                {"detail": "links must be a non-empty list."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        ids = []
        for link in links:
            if not isinstance(link, dict) or link.get("temp_id") is None:
                return Response(
                    {"detail": "Each link needs a temp_id."},
                    status=status.HTTP_400_BAD_REQUEST,
                )
            ids.append(str(link["temp_id"]))
        try:
            queryset = self.get_queryset().filter(pk__in=ids)
        except (ValidationError, ValueError):
            return Response(
                {"detail": "One or more temp_ids are not valid upload ids."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        by_id = {str(temp.id): temp for temp in queryset}

        updated = []
        errors = []
        for link in links:
            temp_id = str(link["temp_id"])
            temp = by_id.get(temp_id)
            if temp is None:
                errors.append({"temp_id": temp_id, "error": "not found"})
                continue
            if temp.status == TempUpload.STATUS_COMPLETED:
                errors.append({"temp_id": temp_id, "error": "already in the library"})
                continue
            trimmed = services.trim_e621_post(link.get("post"))
            post_id = link.get("post_id")
            has_numeric_id = post_id is not None and str(post_id).isdigit()
            if trimmed is None and not has_numeric_id:
                errors.append({"temp_id": temp_id, "error": "no post payload"})
                continue
            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 link.get("auto", True)
                else TempUpload.RESOLUTION_LINKED
            )
            # An MD5 match means the staged file is byte-identical, so there
            # is nothing to download from e621; only link a remote URL when
            # the post's file differs.
            file_data = (trimmed or {}).get("file") or {}
            candidate_url = str(
                link.get("file_url") or file_data.get("url") or ""
            ).strip()
            post_md5 = str(file_data.get("md5") or "").strip().lower()
            if candidate_url and post_md5 and post_md5 == temp.md5.lower():
                candidate_url = ""
            temp.save()
            try:
                complete_temp_upload(temp, download_url=candidate_url or None)
            except Exception as exc:  # noqa: BLE001 - report per-file failures
                logger.exception("Could not complete staged upload %s", temp.id)
                temp.status = TempUpload.STATUS_ERROR
                temp.save(update_fields=["status", "updated_at"])
                errors.append({"temp_id": temp_id, "error": str(exc)})
                continue
            temp.refresh_from_db()
            updated.append(self.get_serializer(temp).data)

        return Response({"updated": updated, "errors": errors})

    @action(detail=False, methods=["post"], url_path="resolve-bulk")
    def resolve_bulk(self, request):
        """Move many staged uploads into the library with one rating.

        The upload board's bulk tool: pick a rating, tick pending files, and
        they are all resolved as custom entries (their staged tags/notes are
        kept). Own rows only, like the rest of the viewset.
        """
        ids = request.data.get("temp_ids")
        if not isinstance(ids, list) or not ids:
            return Response(
                {"detail": "temp_ids must be a non-empty list."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        rating = str(request.data.get("rating") or "").strip()
        if rating not in {"s", "q", "e"}:
            return Response(
                {"detail": "rating must be one of s, q or e."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        try:
            queryset = self.get_queryset().filter(
                pk__in=[str(value) for value in ids]
            )
        except (ValidationError, ValueError):
            return Response(
                {"detail": "One or more ids are not valid upload ids."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        resolved: list[str] = []
        errors: list[dict[str, str]] = []
        for temp in queryset:
            if temp.status == TempUpload.STATUS_COMPLETED:
                errors.append(
                    {"temp_id": str(temp.id), "error": "already in the library"}
                )
                continue
            temp.custom_rating = rating
            temp.resolution = TempUpload.RESOLUTION_CUSTOM
            temp.save(update_fields=["custom_rating", "resolution", "updated_at"])
            try:
                complete_temp_upload(temp)
            except Exception as exc:  # noqa: BLE001 - report per-file failures
                logger.exception("Could not complete staged upload %s", temp.id)
                temp.status = TempUpload.STATUS_ERROR
                temp.save(update_fields=["status", "updated_at"])
                errors.append({"temp_id": str(temp.id), "error": str(exc)})
                continue
            resolved.append(str(temp.id))

        found = {str(temp.id) for temp in queryset}
        for value in sorted({str(value) for value in ids} - found):
            errors.append({"temp_id": value, "error": "not found"})

        return Response({"resolved": resolved, "errors": errors})

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