J621/backend/apps/library/serializers.py

import os
from datetime import timedelta
from pathlib import Path

from django.conf import settings
from rest_framework import serializers

from .models import (
    DownloadTask,
    MatchTask,
    MediaItem,
    MediaLocation,
    SimilarityCheck,
    TempUpload,
)
from .services import (
    MEDIA_FILE_SALT,
    UPLOAD_FILE_SALT,
    VIDEO_EXTENSIONS,
    signed_media_url,
)
from .signing_urls import sign_payload


class MediaLocationSerializer(serializers.ModelSerializer):
    class Meta:
        model = MediaLocation
        fields = ["id", "path", "rel_path", "mtime"]


class MediaItemSerializer(serializers.ModelSerializer):
    tags = serializers.ListField(
        child=serializers.CharField(max_length=100), required=False
    )
    rating = serializers.ChoiceField(choices=["", "s", "q", "e"], required=False)
    locations = MediaLocationSerializer(many=True, read_only=True)
    j_id = serializers.SerializerMethodField()
    display_rating = serializers.SerializerMethodField()
    raw_url = serializers.SerializerMethodField()
    thumbnail_url = serializers.SerializerMethodField()
    filename = serializers.SerializerMethodField()
    extension = serializers.SerializerMethodField()
    kind = serializers.SerializerMethodField()
    uploaded_by = serializers.SerializerMethodField()
    uploaded_by_id = serializers.IntegerField(read_only=True)

    class Meta:
        model = MediaItem
        fields = [
            "id",
            "j_id",
            "md5",
            "size",
            "rating",
            "display_rating",
            "tags",
            "notes",
            "filename",
            "extension",
            "kind",
            "raw_url",
            "thumbnail_url",
            "uploaded_by",
            "uploaded_by_id",
            "e621_post_id",
            "e621_data",
            "e621_match_status",
            "e621_checked_at",
            "locations",
            "created_at",
            "updated_at",
        ]
        read_only_fields = [
            "id",
            "j_id",
            "md5",
            "size",
            "display_rating",
            "filename",
            "extension",
            "kind",
            "raw_url",
            "thumbnail_url",
            "uploaded_by",
            "uploaded_by_id",
            "e621_post_id",
            "e621_data",
            "e621_match_status",
            "e621_checked_at",
            "locations",
            "created_at",
            "updated_at",
        ]

    def _request_user(self):
        request = self.context.get("request")
        user = getattr(request, "user", None)
        if user is None or not getattr(user, "is_authenticated", False):
            return None
        return user

    def get_j_id(self, obj):
        return f"J-{obj.id}"

    def get_raw_url(self, obj):
        return signed_media_url(
            obj, self._request_user(), "raw", request=self.context.get("request")
        )

    def get_thumbnail_url(self, obj):
        return signed_media_url(
            obj,
            self._request_user(),
            "thumbnail",
            request=self.context.get("request"),
        )

    def get_display_rating(self, obj):
        if obj.rating:
            return obj.rating
        data = obj.e621_data or {}
        rating = data.get("rating") if isinstance(data, dict) else None
        return rating if rating in {"s", "q", "e"} else ""

    def get_uploaded_by(self, obj):
        return obj.uploaded_by.username if obj.uploaded_by else None

    def _first_location(self, obj):
        locations = list(obj.locations.all())
        return locations[0] if locations else None

    def get_filename(self, obj):
        location = self._first_location(obj)
        return os.path.basename(location.rel_path) if location else obj.md5

    def get_extension(self, obj):
        return os.path.splitext(self.get_filename(obj))[1].lstrip(".").lower()

    def get_kind(self, obj):
        return "video" if "." + self.get_extension(obj) in VIDEO_EXTENSIONS else "image"


class TempUploadSerializer(serializers.ModelSerializer):
    temp_id = serializers.UUIDField(source="id", read_only=True)
    library_j_id = serializers.SerializerMethodField()
    file_url = serializers.SerializerMethodField()
    preview_url = serializers.SerializerMethodField()
    md5_checked = serializers.SerializerMethodField()
    visual_checked = serializers.SerializerMethodField()
    iqdb_checked = serializers.SerializerMethodField()
    processing = serializers.SerializerMethodField()
    similar_count = serializers.SerializerMethodField()
    visual_matches = serializers.SerializerMethodField()

    class Meta:
        model = TempUpload
        fields = [
            "temp_id",
            "original_filename",
            "md5",
            "size",
            "status",
            "resolution",
            "e621_post_id",
            "e621_data",
            "custom_rating",
            "custom_tags",
            "custom_notes",
            "iqdb_data",
            "visual_matches",
            "library_j_id",
            "file_url",
            "preview_url",
            "pipeline_error",
            "attempts",
            "md5_checked",
            "visual_checked",
            "iqdb_checked",
            "processing",
            "similar_count",
            "created_at",
            "updated_at",
        ]
        read_only_fields = fields

    def _request_user(self):
        request = self.context.get("request")
        user = getattr(request, "user", None)
        if user is None or not getattr(user, "is_authenticated", False):
            return None
        return user

    def get_library_j_id(self, obj):
        return f"J-{obj.library_item_id}" if obj.library_item_id else None

    def get_file_url(self, obj):
        """Signed URL so <img>/<video> tags can fetch the staged file."""
        if not obj.file:
            return None
        user = self._request_user()
        if user is None:
            return None
        signature = sign_payload(
            {"temp": str(obj.id), "user": user.id},
            UPLOAD_FILE_SALT,
        )
        url = f"/api/uploads/{obj.id}/file/?sig={signature}"
        request = self.context.get("request")
        return request.build_absolute_uri(url) if request is not None else url

    def get_preview_url(self, obj):
        """A URL an <img>/<video> tag can load: staged file or library item."""
        user = self._request_user()
        if user is None:
            return None
        if obj.file:
            return self.get_file_url(obj)
        item = obj.library_item
        if item is None:
            return None
        location = item.locations.first()
        is_video = (
            location is not None
            and Path(location.rel_path).suffix.lower() in VIDEO_EXTENSIONS
        )
        action = "thumbnail" if is_video else "raw"
        return signed_media_url(
            item, user, action, request=self.context.get("request")
        )

    def get_md5_checked(self, obj):
        return obj.e621_checked_at is not None

    def get_visual_checked(self, obj):
        return obj.visual_checked_at is not None

    def get_iqdb_checked(self, obj):
        return obj.iqdb_data is not None

    def get_processing(self, obj):
        return obj.claimed_at is not None

    def get_similar_count(self, obj):
        return len(obj.iqdb_data or []) + len(obj.visual_matches or [])

    @staticmethod
    def _visual_item_id(entry):
        if not isinstance(entry, dict):
            return None
        item_id = entry.get("item_id")
        if item_id is None:
            j_id = str(entry.get("j_id") or "")
            if j_id.upper().startswith("J-"):
                j_id = j_id[2:]
            item_id = j_id if j_id.isdigit() else None
        try:
            return int(item_id)
        except (TypeError, ValueError):
            return None

    def get_visual_matches(self, obj):
        """Rebuild match rows with fresh signed thumbnail URLs.

        Storing the signed URL meant it aged out (or came from an older
        signing scheme) and the "Already in your library" grid showed broken
        tiles. The stored rows only carry the item reference now.
        """
        entries = obj.visual_matches or []
        if not entries:
            return entries
        wanted = {}
        for entry in entries:
            item_id = self._visual_item_id(entry)
            if item_id is not None:
                wanted[item_id] = None
        items = MediaItem.objects.in_bulk(list(wanted))
        user = self._request_user()
        request = self.context.get("request")
        matches = []
        for entry in entries:
            item_id = self._visual_item_id(entry)
            item = items.get(item_id) if item_id is not None else None
            if item is None:
                continue
            matches.append(
                {
                    "j_id": f"J-{item.id}",
                    "filename": entry.get("filename") or item.md5,
                    "similarity": entry.get("similarity"),
                    "thumbnail_url": signed_media_url(
                        item, user, "thumbnail", request=request
                    ),
                }
            )
        return matches


class TempUploadListSerializer(TempUploadSerializer):
    """Compact staged-upload row for the board and the status polling.

    Drops the heavy post/IQDB payloads (the metadata modal fetches the full
    row) while keeping the pipeline flags the board renders per tile.
    """

    class Meta(TempUploadSerializer.Meta):
        fields = [
            field
            for field in TempUploadSerializer.Meta.fields
            if field
            not in {"e621_data", "iqdb_data", "visual_matches", "custom_tags", "custom_notes"}
        ]
        read_only_fields = fields


class DownloadTaskSerializer(serializers.ModelSerializer):
    task_id = serializers.UUIDField(source="id", read_only=True)
    library_j_id = serializers.SerializerMethodField()

    class Meta:
        model = DownloadTask
        fields = [
            "task_id",
            "post_id",
            "filename",
            "status",
            "progress",
            "downloaded",
            "total",
            "speed",
            "error",
            "library_j_id",
            "created_at",
            "updated_at",
        ]
        read_only_fields = fields

    def get_library_j_id(self, obj):
        return f"J-{obj.library_item_id}" if obj.library_item_id else None


class MatchTaskSerializer(serializers.ModelSerializer):
    task_id = serializers.UUIDField(source="id", read_only=True)

    class Meta:
        model = MatchTask
        fields = [
            "task_id",
            "scope",
            "status",
            "total",
            "processed",
            "matched",
            "not_found",
            "deleted",
            "failed",
            "error",
            "created_at",
            "updated_at",
        ]
        read_only_fields = fields


class SimilarityCheckSerializer(serializers.ModelSerializer):
    check_id = serializers.UUIDField(source="id", read_only=True)
    file_url = serializers.SerializerMethodField()
    expires_at = serializers.SerializerMethodField()

    class Meta:
        model = SimilarityCheck
        fields = [
            "check_id",
            "original_filename",
            "md5",
            "size",
            "file_url",
            "results",
            "expires_at",
            "created_at",
        ]
        read_only_fields = fields

    def _request_user(self):
        request = self.context.get("request")
        user = getattr(request, "user", None)
        if user is None or not getattr(user, "is_authenticated", False):
            return None
        return user

    def get_file_url(self, obj):
        """Signed URL so <img> tags can preview the temporary upload."""
        user = self._request_user()
        if user is None or not obj.file:
            return None
        signature = sign_payload(
            {"check": str(obj.id), "user": user.id},
            UPLOAD_FILE_SALT,
        )
        url = f"/api/similarity/{obj.id}/file/?sig={signature}"
        request = self.context.get("request")
        return request.build_absolute_uri(url) if request is not None else url

    def get_expires_at(self, obj):
        ttl = int(getattr(settings, "SIMILARITY_TTL_MINUTES", 30))
        return obj.created_at + timedelta(minutes=ttl)