J621/backend/apps/library/views.py

import re
import shutil
import uuid
from pathlib import Path
from urllib.parse import urlparse

from django.conf import settings
from django.core import signing
from django.db.models import Min, Q
from django.http import Http404, StreamingHttpResponse
from django.shortcuts import get_object_or_404
from django.utils import timezone
from django.utils.text import get_valid_filename
from rest_framework import mixins, status, viewsets
from rest_framework.decorators import action
from rest_framework.permissions import AllowAny, IsAuthenticatedOrReadOnly
from rest_framework.response import Response
from rest_framework.throttling import ScopedRateThrottle
from rest_framework.views import APIView

from . import e621, matching, services
from .downloads import reap_stale_downloads, start_download_task
from .matching import reap_stale_match_tasks, start_match_task
from .models import DownloadTask, MatchTask, MediaItem
from .permissions import CanUpload, IsUploaderOrStaffOrReadOnly
from .serializers import (
    DownloadTaskSerializer,
    MatchTaskSerializer,
    MediaItemSerializer,
)

LIST_ORDERINGS = {"name", "-name", "size", "-size", "created_at", "-created_at"}
MD5_RE = re.compile(r"[0-9a-fA-F]{32}")


class MediaItemViewSet(
    mixins.ListModelMixin,
    mixins.RetrieveModelMixin,
    mixins.UpdateModelMixin,
    viewsets.GenericViewSet,
):
    serializer_class = MediaItemSerializer
    permission_classes = [IsAuthenticatedOrReadOnly, IsUploaderOrStaffOrReadOnly]
    http_method_names = ["get", "post", "patch", "head", "options"]

    def get_queryset(self):
        queryset = MediaItem.objects.prefetch_related("locations").annotate(
            name=Min("locations__rel_path")
        )
        if not self.request.user.is_authenticated:
            queryset = queryset.filter(hidden_from_guests=False)
        j_ids = self.request.query_params.get("j_ids", "").strip()
        if j_ids:
            numeric_ids = []
            for value in j_ids.split(","):
                text = value.strip()
                number = text[2:] if text.upper().startswith("J-") else text
                if number.isdigit():
                    numeric_ids.append(int(number))
            queryset = queryset.filter(pk__in=numeric_ids)
        search = self.request.query_params.get("search", "").strip()
        search_type = self.request.query_params.get("search_type", "filename").strip()
        if search:
            words = [word for word in search.lower().split() if word]
            if search_type == "tags":
                for word in words:
                    queryset = queryset.filter(search_tags__icontains=word)
            elif search_type == "both":
                query = Q(locations__rel_path__icontains=search)
                tag_query = Q()
                for word in words:
                    tag_query &= Q(search_tags__icontains=word)
                queryset = queryset.filter(query | tag_query)
            else:
                queryset = queryset.filter(locations__rel_path__icontains=search)
        ratings = [
            value
            for value in self.request.query_params.get("rating", "").split(",")
            if value in {"s", "q", "e"}
        ]
        if ratings:
            queryset = queryset.filter(rating__in=ratings)
        statuses = []
        for value in self.request.query_params.getlist("status"):
            statuses.extend(
                part.strip() for part in value.split(",") if part.strip()
            )
        if statuses:
            status_query = Q()
            for status_value in set(statuses):
                if status_value == "custom":
                    status_query |= Q(has_custom_data=True)
                elif status_value == "matched":
                    status_query |= Q(e621_match_status=MediaItem.E621_MATCHED)
                elif status_value == "not_found":
                    status_query |= Q(e621_match_status=MediaItem.E621_NOT_FOUND)
                elif status_value == "deleted":
                    status_query |= Q(e621_match_status=MediaItem.E621_DELETED)
                elif status_value == "unknown":
                    status_query |= Q(
                        has_custom_data=False,
                        e621_match_status=MediaItem.E621_UNKNOWN,
                    )
            queryset = queryset.filter(status_query)
        ordering = self.request.query_params.get("ordering", "").strip()
        queryset = queryset.order_by(
            ordering if ordering in LIST_ORDERINGS else "-created_at"
        )
        return queryset.distinct()

    def get_object(self):
        """Resolve J-<id>, a bare numeric id, or an MD5 fingerprint."""
        value = str(self.kwargs.get("pk", "")).strip()
        queryset = self.get_queryset()
        if MD5_RE.fullmatch(value):
            obj = get_object_or_404(queryset, md5=value.lower())
        else:
            numeric = value[2:] if value.upper().startswith("J-") else value
            if not numeric.isdigit():
                raise Http404
            obj = get_object_or_404(queryset, pk=int(numeric))
        self.check_object_permissions(self.request, obj)
        return obj

    def _signed_media_item(self, request, action_name):
        """Signed media URLs let <img>/<video> tags bypass guest filtering."""
        if request.user.is_authenticated:
            return None
        signature = request.query_params.get("sig")
        if not signature:
            return None
        try:
            payload = signing.loads(
                signature, salt=services.MEDIA_FILE_SALT, max_age=86400
            )
        except signing.BadSignature:
            return None
        if payload.get("action") != action_name:
            return None
        return MediaItem.objects.filter(pk=payload.get("item")).first()

    def _media_object(self, request, action_name):
        item = self._signed_media_item(request, action_name)
        if item is not None:
            return item
        return self.get_object()

    @action(detail=True, methods=["get"])
    def raw(self, request, pk=None):
        item = self._media_object(request, "raw")
        location = item.locations.first()
        if location is None:
            return Response(
                {"detail": "No file for this item."},
                status=status.HTTP_404_NOT_FOUND,
            )
        return services.serve_file(
            request, location.path, download=request.query_params.get("download") == "1"
        )

    @action(detail=True, methods=["get"])
    def thumbnail(self, request, pk=None):
        item = self._media_object(request, "thumbnail")
        location = item.locations.first()
        if location is None:
            return Response(
                {"detail": "No file for this item."},
                status=status.HTTP_404_NOT_FOUND,
            )
        path = Path(location.path)
        if path.suffix.lower() in services.VIDEO_EXTENSIONS:
            thumbnail = services.generate_video_thumbnail(item.md5, path)
            if thumbnail is None:
                return Response(
                    {"detail": "Thumbnail unavailable."},
                    status=status.HTTP_404_NOT_FOUND,
                )
            return services.serve_file(request, thumbnail)
        return services.serve_file(request, path)

    @action(detail=False, methods=["post"], permission_classes=[AllowAny])
    def lookup(self, request):
        """Report which of the given MD5s exist in the library."""
        md5s = request.data.get("md5s")
        if not isinstance(md5s, list):
            return Response(
                {"detail": "md5s must be a list."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        wanted = {str(value).strip().lower()[:32] for value in md5s if value}
        queryset = MediaItem.objects.filter(md5__in=wanted)
        if not request.user.is_authenticated:
            queryset = queryset.filter(hidden_from_guests=False)
        rows = list(queryset.values_list("md5", "id"))
        return Response(
            {
                "found": sorted(md5 for md5, _ in rows),
                "j_ids": {md5: f"J-{item_id}" for md5, item_id in rows},
            }
        )

    @action(detail=True, methods=["get"])
    def neighbors(self, request, pk=None):
        """Previous/next items in the ordered list, for keyboard navigation."""
        item = self.get_object()
        ordering = request.query_params.get("ordering", "-created_at").strip()
        if ordering not in LIST_ORDERINGS:
            ordering = "-created_at"
        field = ordering.lstrip("-")
        descending = ordering.startswith("-")
        value = {
            "size": item.size,
            "created_at": item.created_at,
            "name": item.name,
        }[field]

        queryset = self.get_queryset().exclude(pk=item.pk)
        if descending:
            previous_query = queryset.filter(**{f"{field}__gt": value}).order_by(field)
            next_query = queryset.filter(**{f"{field}__lt": value}).order_by(f"-{field}")
        else:
            previous_query = queryset.filter(**{f"{field}__lt": value}).order_by(
                f"-{field}"
            )
            next_query = queryset.filter(**{f"{field}__gt": value}).order_by(field)

        def brief(obj):
            location = obj.locations.first()
            return {
                "j_id": f"J-{obj.id}",
                "md5": obj.md5,
                "filename": Path(location.rel_path).name if location else obj.md5,
            }

        previous = previous_query.first()
        next_item = next_query.first()
        return Response(
            {
                "previous": brief(previous) if previous else None,
                "next": brief(next_item) if next_item else None,
            }
        )

    def _can_match(self, request, item):
        user = request.user
        return bool(user.is_app_staff or item.uploaded_by_id == user.id)

    @action(detail=True, methods=["post"], permission_classes=[CanUpload])
    def match(self, request, pk=None):
        """Link this item to e621: MD5 lookup, or an explicit post id."""
        item = self.get_object()
        if not self._can_match(request, item):
            return Response(
                {"detail": "Only the uploader or staff can match this item."},
                status=status.HTTP_403_FORBIDDEN,
            )
        extra = {}
        post_id = request.data.get("post_id")
        try:
            if post_id not in (None, ""):
                text = str(post_id).strip()
                if not text.isdigit():
                    return Response(
                        {"detail": "Post ID must be a number."},
                        status=status.HTTP_400_BAD_REQUEST,
                    )
                extra["md5_match"] = matching.link_item_to_post(
                    item, int(text), request.user
                )
            else:
                matching.check_item(item, request.user)
        except e621.E621NotFound:
            return Response(
                {"detail": "That post does not exist on e621."},
                status=status.HTTP_404_NOT_FOUND,
            )
        except e621.E621Error as exc:
            return Response(
                {"detail": str(exc)}, status=status.HTTP_502_BAD_GATEWAY
            )
        item.refresh_from_db()
        data = dict(self.get_serializer(item).data)
        data.update(extra)
        return Response(data)

    @action(detail=True, methods=["post"], permission_classes=[CanUpload])
    def unlink(self, request, pk=None):
        """Forget this item's e621 link so it can be matched again."""
        item = self.get_object()
        if not self._can_match(request, item):
            return Response(
                {"detail": "Only the uploader or staff can unlink this item."},
                status=status.HTTP_403_FORBIDDEN,
            )
        matching.unlink_item(item)
        return Response(self.get_serializer(item).data)

    @action(detail=True, methods=["post"], permission_classes=[CanUpload])
    def optimize(self, request, pk=None):
        """Replace this item's file(s) with a client-processed version.

        The browser does the actual compression (WebCodecs/WASM); this only
        applies the result. The item keeps its J-ID while MD5, size and
        perceptual hashes are recomputed, so it may become a duplicate again.
        """
        item = self.get_object()
        if not self._can_match(request, item):
            return Response(
                {"detail": "Only the uploader or staff can optimize this item."},
                status=status.HTTP_403_FORBIDDEN,
            )
        upload = request.FILES.get("file")
        if upload is None:
            return Response(
                {"detail": "A processed 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,
            )
        locations = list(item.locations.all())
        if not locations:
            return Response(
                {"detail": "This item has no file on disk."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        folder = Path(locations[0].path).parent
        temp_path = folder / f".optimize-{uuid.uuid4().hex}{extension}"
        try:
            with open(temp_path, "wb") as handle:
                for chunk in upload.chunks():
                    handle.write(chunk)
            new_md5 = services.compute_md5(temp_path)
            new_size = temp_path.stat().st_size
        except OSError as exc:
            temp_path.unlink(missing_ok=True)
            return Response(
                {"detail": f"Could not store the processed file: {exc}"},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )

        if new_md5 == item.md5:
            temp_path.unlink(missing_ok=True)
            return Response(
                {"detail": "The processed file is identical to the current one."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        clash = MediaItem.objects.filter(md5=new_md5).exclude(pk=item.pk).first()
        if clash is not None:
            temp_path.unlink(missing_ok=True)
            return Response(
                {
                    "detail": (
                        f"The processed file is an exact match (MD5) for "
                        f"J-{clash.id}; nothing was changed."
                    )
                },
                status=status.HTTP_409_CONFLICT,
            )

        new_name = f"J-{item.id}{extension}"
        try:
            for location in locations:
                target = Path(location.path).parent / new_name
                shutil.copyfile(temp_path, target)
                if str(target) != location.path:
                    Path(location.path).unlink(missing_ok=True)
                    location.path = str(target)
                    try:
                        location.rel_path = str(
                            target.relative_to(settings.WATCHED_FOLDER)
                        )
                    except ValueError:
                        location.rel_path = target.name
                location.mtime = target.stat().st_mtime
                location.save(update_fields=["path", "rel_path", "mtime"])
        except OSError as exc:
            return Response(
                {"detail": f"Could not replace the file(s): {exc}"},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )
        finally:
            temp_path.unlink(missing_ok=True)

        hashes = (
            services.compute_visual_hashes(locations[0].path)
            if extension in services.IMAGE_EXTENSIONS
            else {}
        )
        item.md5 = new_md5
        item.size = new_size
        for field in services.HASH_FIELDS:
            setattr(item, field, hashes.get(field, ""))
        item.save(
            update_fields=["md5", "size", *services.HASH_FIELDS, "updated_at"]
        )

        item = self.get_queryset().get(pk=item.pk)
        return Response(self.get_serializer(item).data)


class DownloadTaskViewSet(
    mixins.ListModelMixin,
    mixins.RetrieveModelMixin,
    viewsets.GenericViewSet,
):
    """Background 'Download to Library' jobs."""

    serializer_class = DownloadTaskSerializer
    permission_classes = [CanUpload]
    http_method_names = ["get", "post", "head", "options"]

    def get_queryset(self):
        reap_stale_downloads()
        queryset = DownloadTask.objects.select_related("library_item")
        post_id = self.request.query_params.get("post_id", "").strip()
        if post_id.isdigit():
            queryset = queryset.filter(post_id=int(post_id))
        user = self.request.user
        if not user.is_app_staff:
            queryset = queryset.filter(user=user)
        return queryset

    def create(self, request):
        url = str(request.data.get("url") or "").strip()
        post_id = request.data.get("post_id")
        filename = str(request.data.get("filename") or "").strip()
        try:
            services.validate_remote_url(url)
        except services.RemoteUrlError as exc:
            return Response(
                {"detail": str(exc)},
                status=status.HTTP_400_BAD_REQUEST,
            )
        trimmed = services.trim_e621_post(request.data.get("post"))
        resolved_post_id = None
        if trimmed is not None and trimmed.get("id") is not None:
            resolved_post_id = int(trimmed["id"])
        elif post_id is not None and str(post_id).isdigit():
            resolved_post_id = int(post_id)

        task = DownloadTask.objects.create(
            user=request.user,
            post_id=resolved_post_id,
            url=url,
            filename=filename,
            e621_data=trimmed,
        )
        start_download_task(task.id)
        return Response(
            self.get_serializer(task).data, status=status.HTTP_202_ACCEPTED
        )

    @action(detail=True, methods=["post"])
    def cancel(self, request, pk=None):
        task = self.get_object()
        if task.status not in {
            DownloadTask.STATUS_PENDING,
            DownloadTask.STATUS_DOWNLOADING,
        }:
            return Response(
                {"detail": "Task is not running."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        DownloadTask.objects.filter(id=task.id).update(
            cancelled=True,
            status=DownloadTask.STATUS_CANCELLED,
            speed=None,
            updated_at=timezone.now(),
        )
        return Response({"success": True})


class MatchTaskViewSet(
    mixins.ListModelMixin,
    mixins.RetrieveModelMixin,
    viewsets.GenericViewSet,
):
    """Background library-wide e621 match scans."""

    serializer_class = MatchTaskSerializer
    permission_classes = [CanUpload]
    http_method_names = ["get", "post", "head", "options"]

    def get_queryset(self):
        reap_stale_match_tasks()
        queryset = MatchTask.objects.all()
        user = self.request.user
        if not user.is_app_staff:
            queryset = queryset.filter(user=user)
        return queryset

    def create(self, request):
        scope = str(request.data.get("scope") or MatchTask.SCOPE_MISSING).strip()
        if scope not in {MatchTask.SCOPE_MISSING, MatchTask.SCOPE_ALL}:
            scope = MatchTask.SCOPE_MISSING
        task = MatchTask.objects.create(user=request.user, scope=scope)
        start_match_task(task.id)
        return Response(
            self.get_serializer(task).data, status=status.HTTP_202_ACCEPTED
        )

    @action(detail=True, methods=["post"])
    def cancel(self, request, pk=None):
        task = self.get_object()
        if task.status not in {MatchTask.STATUS_PENDING, MatchTask.STATUS_RUNNING}:
            return Response(
                {"detail": "Task is not running."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        MatchTask.objects.filter(id=task.id).update(
            cancelled=True,
            status=MatchTask.STATUS_CANCELLED,
            updated_at=timezone.now(),
        )
        return Response({"success": True})


class ClientDownloadView(APIView):
    """Stream an e621 file straight to the browser (no library write)."""

    permission_classes = [AllowAny]
    throttle_classes = [ScopedRateThrottle]
    throttle_scope = "e621_proxy"

    def get(self, request):
        url = str(request.query_params.get("url") or "").strip()
        filename = str(request.query_params.get("filename") or "").strip()
        try:
            services.validate_remote_url(url)
        except services.RemoteUrlError:
            return Response(
                {"detail": "URL not allowed."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        import requests

        try:
            upstream = services.open_remote(
                url,
                headers={"User-Agent": settings.USER_AGENT},
                stream=True,
                timeout=60,
            )
            upstream.raise_for_status()
        except services.RemoteUrlError as exc:
            return Response(
                {"detail": str(exc)}, status=status.HTTP_400_BAD_REQUEST
            )
        except requests.RequestException as exc:
            return Response(
                {"detail": f"Could not fetch the file: {exc}"},
                status=status.HTTP_502_BAD_GATEWAY,
            )

        content_type = upstream.headers.get(
            "Content-Type", "application/octet-stream"
        )
        name = get_valid_filename(
            filename or Path(parsed.path).name or "download"
        )

        def stream():
            try:
                for chunk in upstream.iter_content(chunk_size=64 * 1024):
                    if chunk:
                        yield chunk
            finally:
                upstream.close()

        response = StreamingHttpResponse(stream(), content_type=content_type)
        response["Content-Disposition"] = f'attachment; filename="{name}"'
        length = upstream.headers.get("Content-Length")
        if length:
            response["Content-Length"] = length
        return response