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.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 TempUploadListSerializer, TempUploadSerializer
from .signing_urls import load_payload
from .tools import HASH_FIELDS, hashed_items, hashes_similarity

logger = logging.getLogger(__name__)


def build_hash_index():
    """Library hash mappings for similarity scans, loaded once per batch.

    Only items that actually carry perceptual hashes are included; the old
    per-file scan walked every row (including videos and unchecked items).
    """
    algorithms = list(HASH_FIELDS)
    return [
        (item, {field: getattr(item, field, "") for field in algorithms})
        for item in hashed_items(algorithms)
    ]


def match_hashes(hashes, index, limit=10, user=None, request=None):
    """Library items whose perceptual hashes are close to ``hashes``."""
    algorithms = list(HASH_FIELDS)
    threshold = settings.VISUAL_MATCH_THRESHOLD
    matches = []
    for item, item_hashes in index:
        similarity = hashes_similarity(hashes, item_hashes, 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 find_library_matches(path, limit=10, user=None, request=None, index=None):
    """Library items visually similar to a staged file.

    Pass a prebuilt ``index`` (see build_hash_index) to reuse it across a
    whole batch instead of rescanning the library per file.
    """
    hashes = services.compute_visual_hashes(path)
    if not hashes:
        return []
    if index is None:
        index = build_hash_index()
    return match_hashes(hashes, index, limit=limit, user=user, request=request)


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()
    from .upload_pipeline import record_completion

    record_completion(temp, item)
    # The board learns about completions from the live feed, so the record is
    # deleted as soon as the file is indexed: nothing left to dismiss. Delete
    # through the queryset so callers keep ``temp.pk`` for their response.
    TempUpload.objects.filter(pk=temp.pk).delete()
    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_serializer_class(self):
        # The board polls the list, so its payload stays small; the metadata
        # modal fetches the full row from the detail endpoint.
        if self.action == "list":
            return TempUploadListSerializer
        return TempUploadSerializer

    def get_queryset(self):
        # Staged uploads are private: everyone, staff included, only sees
        # their own board. (The file action still lets staff read bytes by id
        # for support purposes.)
        return TempUpload.objects.select_related("library_item").filter(
            user=self.request.user
        )

    def _completed_payload(self, temp, item):
        """Synthetic row for an upload that is indexed immediately.

        Duplicates and resolved uploads never leave a board record; the SPA
        turns this response (or the live completion feed) into a "J-x
        uploaded" card that lives only in the page session.
        """
        return {
            "temp_id": str(temp.pk),
            "original_filename": temp.original_filename,
            "md5": temp.md5,
            "size": temp.size,
            "status": TempUpload.STATUS_COMPLETED,
            "resolution": temp.resolution,
            "e621_post_id": temp.e621_post_id,
            "library_j_id": f"J-{item.id}",
            "file_url": None,
            "preview_url": services.signed_media_url(
                item, self.request.user, "thumbnail", request=self.request
            ),
        }

    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()
            from .upload_pipeline import record_completion

            record_completion(temp, existing)
            payload = self._completed_payload(temp, existing)
            TempUpload.objects.filter(pk=temp.pk).delete()
            return Response(payload, status=status.HTTP_201_CREATED)
        # Visual similarity and IQDB run in the background pipeline so a large
        # batch uploads at full speed and the work survives the browser.
        temp.save()
        # Kick the server-side pipeline; staging no longer waits on e621 and
        # the work continues even if the browser navigates away.
        from .upload_pipeline import start_pipeline

        start_pipeline(request.user)
        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=False, methods=["get"])
    def status(self, request):
        """Cheap pipeline state for the shell indicator and the upload page."""
        from .upload_pipeline import status_payload

        return Response(status_payload(request.user, request=request))

    @action(detail=False, methods=["post"])
    def process(self, request):
        """Start (or resume) the pipeline for the caller's staged uploads.

        Idempotent: the client calls this after staging files, on page load
        and when a paused run should be retried.
        """
        from .upload_pipeline import start_pipeline, status_payload

        start_pipeline(request.user)
        return Response(status_payload(request.user, request=request))

    @action(detail=True, methods=["post"])
    def retry(self, request, pk=None):
        """Queue one staged upload for another pipeline pass."""
        from .upload_pipeline import start_pipeline, status_payload

        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,
            )
        update = {
            "claimed_at": None,
            "attempts": 0,
            "pipeline_error": "",
            "updated_at": timezone.now(),
        }
        # An explicit phase re-runs that one check even if it already ran.
        phase = str(request.data.get("phase") or "").strip()
        if phase == "md5":
            update["e621_checked_at"] = None
        elif phase == "visual":
            update["visual_matches"] = None
            update["visual_checked_at"] = None
        elif phase == "iqdb":
            update["iqdb_data"] = None
        if temp.status == TempUpload.STATUS_ERROR:
            # A failed import has to go through the MD5 phase again so the
            # completion is retried; other errors only re-run missing phases.
            update["status"] = TempUpload.STATUS_PENDING
            update["e621_checked_at"] = None
        elif temp.status not in (
            TempUpload.STATUS_PENDING,
            TempUpload.STATUS_VISUAL_MATCH,
        ):
            update["status"] = TempUpload.STATUS_PENDING
        TempUpload.objects.filter(pk=temp.pk).update(**update)
        start_pipeline(request.user)
        return Response(status_payload(request.user, request=request))

    @action(detail=False, methods=["post"], url_path="retry-all")
    def retry_all(self, request):
        """Queue every retryable staged upload for another pipeline pass."""
        from .upload_pipeline import start_pipeline, status_payload

        now = timezone.now()
        retryable = self.get_queryset().filter(
            status__in=[
                TempUpload.STATUS_PENDING,
                TempUpload.STATUS_VISUAL_MATCH,
                TempUpload.STATUS_ERROR,
            ]
        )
        retryable.exclude(file="").update(
            claimed_at=None,
            attempts=0,
            pipeline_error="",
            updated_at=now,
        )
        retryable.exclude(file="").filter(status=TempUpload.STATUS_ERROR).update(
            status=TempUpload.STATUS_PENDING,
            e621_checked_at=None,
        )
        start_pipeline(request.user)
        return Response(status_payload(request.user, request=request))

    @action(detail=False, methods=["post"], url_path="discard-bulk")
    def discard_bulk(self, request):
        """Discard many staged uploads in one request (the board's "all")."""
        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,
            )
        if len(ids) > 1000:
            return Response(
                {"detail": "Too many ids in one request (max 1000)."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        values = list(dict.fromkeys(str(value) for value in ids))
        try:
            queryset = self.get_queryset().filter(pk__in=values)
        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.pk): temp for temp in queryset}

        discarded: list[str] = []
        errors: list[dict[str, str]] = []
        for value in values:
            temp = by_id.get(value)
            if temp is None:
                errors.append({"temp_id": value, "error": "not found"})
                continue
            try:
                self.perform_destroy(temp)
                discarded.append(value)
            except Exception as exc:  # noqa: BLE001 - report per-file failures
                logger.exception("Could not discard staged upload %s", value)
                errors.append({"temp_id": value, "error": str(exc)})
        return Response({"discarded": discarded, "errors": errors})

    @action(
        detail=True,
        methods=["get", "head"],
        permission_classes=[AllowAny],
        throttle_classes=[],
    )
    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:
                payload = load_payload(signature, services.UPLOAD_FILE_SALT)
                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:
            item = 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,
            )
        return Response(self._completed_payload(temp, item))

    @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:
                item = 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
            updated.append(self._completed_payload(temp, item))

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