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

    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 [])


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)