J621/backend/apps/library/models.py

from django.conf import settings
from django.db import models

import uuid


class MediaItem(models.Model):
    """A logical media file, identified by its MD5 fingerprint."""

    md5 = models.CharField(max_length=32, unique=True)
    size = models.BigIntegerField(default=0)
    # Custom rating: "s" (safe), "q" (questionable), "e" (explicit) or "" (unrated).
    rating = models.CharField(max_length=1, blank=True, default="")
    tags = models.JSONField(default=list, blank=True)
    notes = models.TextField(blank=True, default="")
    uploaded_by = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="uploads",
    )
    # Guests never see items whose tags hit the mirrored default blacklist.
    hidden_from_guests = models.BooleanField(default=False, db_index=True)
    e621_post_id = models.IntegerField(null=True, blank=True, db_index=True)
    e621_data = models.JSONField(null=True, blank=True)
    # Perceptual hashes (hex strings) used by the duplicates engine.
    ahash = models.CharField(max_length=32, blank=True, default="", db_index=True)
    dhash = models.CharField(max_length=32, blank=True, default="", db_index=True)
    phash = models.CharField(max_length=32, blank=True, default="", db_index=True)
    whash = models.CharField(max_length=32, blank=True, default="", db_index=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ["-created_at"]

    def __str__(self):
        return f"J-{self.pk} ({self.md5})"

    def save(self, *args, **kwargs):
        from .guest_filter import item_is_hidden_for_guests

        self.hidden_from_guests = item_is_hidden_for_guests(self)
        update_fields = kwargs.get("update_fields")
        if update_fields is not None:
            kwargs["update_fields"] = set(update_fields) | {"hidden_from_guests"}
        super().save(*args, **kwargs)


class MediaLocation(models.Model):
    """A physical copy of a media item inside the watched folder."""

    item = models.ForeignKey(
        MediaItem, related_name="locations", on_delete=models.CASCADE
    )
    path = models.CharField(max_length=1024)
    rel_path = models.CharField(max_length=1024)
    mtime = models.FloatField(default=0)

    class Meta:
        unique_together = ("item", "path")
        ordering = ["rel_path"]

    def __str__(self):
        return self.rel_path


class TempUpload(models.Model):
    """A staged upload that has not been accepted into the library yet."""

    STATUS_PENDING = "pending"
    STATUS_VISUAL_MATCH = "visual_match"
    STATUS_COMPLETED = "completed"
    STATUS_ERROR = "error"
    STATUS_CHOICES = [
        (STATUS_PENDING, "Pending"),
        (STATUS_VISUAL_MATCH, "Visual Similarity Detected"),
        (STATUS_COMPLETED, "Completed"),
        (STATUS_ERROR, "Error"),
    ]

    RESOLUTION_AUTO_MD5 = "auto_md5"
    RESOLUTION_DUPLICATE = "duplicate"
    RESOLUTION_LINKED = "linked"
    RESOLUTION_CUSTOM = "custom"

    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    user = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        on_delete=models.CASCADE,
        related_name="temp_uploads",
    )
    file = models.FileField(upload_to="uploads/temp/", blank=True)
    original_filename = models.CharField(max_length=255)
    md5 = models.CharField(max_length=32, db_index=True)
    size = models.BigIntegerField(default=0)
    status = models.CharField(
        max_length=20, choices=STATUS_CHOICES, default=STATUS_PENDING
    )
    resolution = models.CharField(max_length=20, blank=True, default="")
    e621_post_id = models.IntegerField(null=True, blank=True)
    e621_data = models.JSONField(null=True, blank=True)
    custom_rating = models.CharField(max_length=1, blank=True, default="")
    custom_tags = models.JSONField(default=list, blank=True)
    custom_notes = models.TextField(blank=True, default="")
    iqdb_data = models.JSONField(null=True, blank=True)
    visual_matches = models.JSONField(null=True, blank=True)
    library_item = models.ForeignKey(
        MediaItem,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="temp_uploads",
    )
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ["-created_at"]

    def __str__(self):
        return f"{self.original_filename} ({self.status})"


class DownloadTask(models.Model):
    """A background 'Download to Library' job with progress tracking."""

    STATUS_PENDING = "pending"
    STATUS_DOWNLOADING = "downloading"
    STATUS_COMPLETE = "complete"
    STATUS_ERROR = "error"
    STATUS_CANCELLED = "cancelled"
    STATUS_CHOICES = [
        (STATUS_PENDING, "Pending"),
        (STATUS_DOWNLOADING, "Downloading"),
        (STATUS_COMPLETE, "Complete"),
        (STATUS_ERROR, "Error"),
        (STATUS_CANCELLED, "Cancelled"),
    ]

    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    user = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="download_tasks",
    )
    post_id = models.IntegerField(null=True, blank=True)
    url = models.CharField(max_length=1000)
    filename = models.CharField(max_length=255, blank=True, default="")
    e621_data = models.JSONField(null=True, blank=True)
    status = models.CharField(
        max_length=20, choices=STATUS_CHOICES, default=STATUS_PENDING
    )
    progress = models.IntegerField(default=0)
    downloaded = models.BigIntegerField(default=0)
    total = models.BigIntegerField(default=0)
    speed = models.FloatField(null=True, blank=True)
    error = models.TextField(blank=True, default="")
    cancelled = models.BooleanField(default=False)
    library_item = models.ForeignKey(
        MediaItem,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name="download_tasks",
    )
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ["-created_at"]

    def __str__(self):
        return f"Download {self.id} ({self.status})"