J621/backend/apps/library/views.py

import re
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.views import APIView

from . import services
from .downloads import reap_stale_downloads, start_download_task
from .models import DownloadTask, MediaItem
from .permissions import CanUpload, IsUploaderOrStaffOrReadOnly
from .serializers import DownloadTaskSerializer, 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(
                        has_custom_data=False, e621_post_id__isnull=False
                    )
                elif status_value == "unknown":
                    status_query |= Q(
                        has_custom_data=False, e621_post_id__isnull=True
                    )
            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,
            }
        )


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")
        user = self.request.user
        if not (user.is_staff or user.is_superuser):
            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()
        if not url.startswith(("http://", "https://")):
            return Response(
                {"detail": "A valid file URL is required."},
                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, 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]

    def get(self, request):
        url = str(request.query_params.get("url") or "").strip()
        filename = str(request.query_params.get("filename") or "").strip()
        parsed = urlparse(url)
        if (
            parsed.scheme not in {"http", "https"}
            or parsed.hostname not in settings.E621_MEDIA_HOSTS
        ):
            return Response(
                {"detail": "URL not allowed."},
                status=status.HTTP_400_BAD_REQUEST,
            )

        import requests

        try:
            upstream = requests.get(
                url,
                headers={"User-Agent": settings.USER_AGENT},
                stream=True,
                timeout=60,
            )
            upstream.raise_for_status()
        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