J621/backend/apps/library/models.py

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

import uuid


TAG_CLOUD_CACHE_KEYS = [
    "j621.library.tag_cloud.auth",
    "j621.library.tag_cloud.guest",
]


def build_search_tags(tags, e621_data):
    """Lowercase tag text (custom + e621) used for library tag search."""
    names = [str(tag) for tag in (tags or [])]
    data = e621_data if isinstance(e621_data, dict) else {}
    categories = data.get("tags")
    if isinstance(categories, dict):
        for values in categories.values():
            if isinstance(values, list):
                names.extend(str(tag) for tag in values)
    return " ".join(
        sorted({str(name).strip().lower() for name in names if str(name).strip()})
    )


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

    E621_UNKNOWN = "unknown"
    E621_MATCHED = "matched"
    E621_NOT_FOUND = "not_found"
    E621_DELETED = "deleted"
    E621_STATUS_CHOICES = [
        (E621_UNKNOWN, "Not checked"),
        (E621_MATCHED, "Matched"),
        (E621_NOT_FOUND, "Not found on e621"),
        (E621_DELETED, "Post deleted on e621"),
    ]

    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)
    # Cached result of the last MD5/manual match against e621.
    e621_match_status = models.CharField(
        max_length=20,
        choices=E621_STATUS_CHOICES,
        default=E621_UNKNOWN,
        db_index=True,
    )
    e621_checked_at = models.DateTimeField(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)
    # Derived search/filter helpers: flattened tag text and custom-data flag.
    search_tags = models.TextField(blank=True, default="")
    has_custom_data = models.BooleanField(default=False, 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 django.core.cache import cache

        from .guest_filter import item_is_hidden_for_guests

        self.hidden_from_guests = item_is_hidden_for_guests(self)
        self.has_custom_data = bool(self.tags or self.notes)
        self.search_tags = build_search_tags(self.tags, self.e621_data)
        update_fields = kwargs.get("update_fields")
        if update_fields is not None:
            kwargs["update_fields"] = set(update_fields) | {
                "hidden_from_guests",
                "has_custom_data",
                "search_tags",
            }
        super().save(*args, **kwargs)
        cache.delete_many(TAG_CLOUD_CACHE_KEYS)


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",
    )
    # Background pipeline bookkeeping (see apps/library/upload_pipeline.py).
    # e621_checked_at/visual_checked_at are set once the corresponding phase
    # ran, so "never checked" and "checked, nothing found" stay distinct.
    e621_checked_at = models.DateTimeField(null=True, blank=True)
    visual_checked_at = models.DateTimeField(null=True, blank=True)
    # Worker claim for cross-process mutual exclusion; stale claims are
    # reaped and the row queued again.
    claimed_at = models.DateTimeField(null=True, blank=True, db_index=True)
    pipeline_error = models.TextField(blank=True, default="")
    attempts = models.PositiveSmallIntegerField(default=0)
    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 UploadRun(models.Model):
    """Per-user state of the background upload pipeline.

    One row per user acts as the cheap status source the SPA polls and as a
    place for batch-level failures (broken e621 credentials, outages) that
    would otherwise be repeated on every row.
    """

    STATUS_IDLE = "idle"
    STATUS_RUNNING = "running"
    STATUS_PAUSED = "paused"
    STATUS_ERROR = "error"
    STATUS_CHOICES = [
        (STATUS_IDLE, "Idle"),
        (STATUS_RUNNING, "Running"),
        (STATUS_PAUSED, "Paused"),
        (STATUS_ERROR, "Error"),
    ]

    PHASE_MD5 = "md5"
    PHASE_VISUAL = "visual"
    PHASE_IQDB = "iqdb"
    PHASE_CHOICES = [
        ("", "None"),
        (PHASE_MD5, "e621 MD5"),
        (PHASE_VISUAL, "Visual similarity"),
        (PHASE_IQDB, "IQDB"),
    ]

    user = models.OneToOneField(
        settings.AUTH_USER_MODEL,
        on_delete=models.CASCADE,
        related_name="upload_run",
    )
    status = models.CharField(
        max_length=20, choices=STATUS_CHOICES, default=STATUS_IDLE
    )
    phase = models.CharField(
        max_length=20, choices=PHASE_CHOICES, blank=True, default=""
    )
    total = models.IntegerField(default=0)
    processed = models.IntegerField(default=0)
    matched = models.IntegerField(default=0)
    failed = models.IntegerField(default=0)
    error = models.TextField(blank=True, default="")
    started_at = models.DateTimeField(null=True, blank=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ["-updated_at"]

    def __str__(self):
        return f"Upload run for {self.user_id} ({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})"


class MatchTask(models.Model):
    """A background library-wide e621 match scan with progress tracking."""

    STATUS_PENDING = "pending"
    STATUS_RUNNING = "running"
    STATUS_COMPLETE = "complete"
    STATUS_ERROR = "error"
    STATUS_CANCELLED = "cancelled"
    STATUS_CHOICES = [
        (STATUS_PENDING, "Pending"),
        (STATUS_RUNNING, "Running"),
        (STATUS_COMPLETE, "Complete"),
        (STATUS_ERROR, "Error"),
        (STATUS_CANCELLED, "Cancelled"),
    ]

    SCOPE_MISSING = "missing"
    SCOPE_ALL = "all"
    SCOPE_CHOICES = [
        (SCOPE_MISSING, "Unchecked items"),
        (SCOPE_ALL, "Whole library"),
    ]

    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="match_tasks",
    )
    scope = models.CharField(
        max_length=20, choices=SCOPE_CHOICES, default=SCOPE_MISSING
    )
    status = models.CharField(
        max_length=20, choices=STATUS_CHOICES, default=STATUS_PENDING
    )
    total = models.IntegerField(default=0)
    processed = models.IntegerField(default=0)
    matched = models.IntegerField(default=0)
    not_found = models.IntegerField(default=0)
    deleted = models.IntegerField(default=0)
    failed = models.IntegerField(default=0)
    error = models.TextField(blank=True, default="")
    cancelled = models.BooleanField(default=False)
    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"Match scan {self.id} ({self.status})"


class SimilarityCheck(models.Model):
    """An ephemeral image uploaded only to check what it looks like.

    Never enters the library: files live in the media temp folder and are
    purged on startup, by the TTL and by `manage.py cleanup_similarity`.
    """

    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="similarity_checks",
    )
    file = models.FileField(upload_to="similarity/", blank=True)
    original_filename = models.CharField(max_length=255)
    md5 = models.CharField(max_length=32, db_index=True, blank=True, default="")
    size = models.BigIntegerField(default=0)
    results = models.JSONField(default=dict, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)

    class Meta:
        ordering = ["-created_at"]

    def __str__(self):
        return f"Similarity check {self.id} ({self.original_filename})"