J621/backend/apps/library/services.py

import hashlib
import json
import logging
import mimetypes
import os
import re
import shutil
import subprocess
from pathlib import Path

import imagehash
from django.conf import settings
from django.http import FileResponse, Http404, HttpResponse
from django.utils.text import get_valid_filename
from PIL import Image

from .models import MediaItem, MediaLocation

logger = logging.getLogger(__name__)

HASH_FIELDS = ("ahash", "dhash", "phash", "whash")
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".apng", ".webp"}

ALLOWED_EXTENSIONS = {
    ".jpg",
    ".jpeg",
    ".png",
    ".gif",
    ".apng",
    ".webp",
    ".mp4",
    ".webm",
}
VIDEO_EXTENSIONS = {".mp4", ".webm"}
UPLOAD_FILE_SALT = "j621.upload-file"
MEDIA_FILE_SALT = "j621.media-file"
CHUNK_SIZE = 1024 * 1024
RANGE_RE = re.compile(r"bytes=(\d*)-(\d*)$")


def compute_md5(path):
    digest = hashlib.md5()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(CHUNK_SIZE), b""):
            digest.update(chunk)
    return digest.hexdigest()


def index_file(path, folder):
    """Index one file into MediaItem/MediaLocation.

    Returns (item, created_item, location, created_location).
    """
    path = Path(path)
    folder = Path(folder)
    stat = path.stat()
    md5 = compute_md5(path)
    item, created_item = MediaItem.objects.get_or_create(
        md5=md5, defaults={"size": stat.st_size}
    )
    if not created_item and item.size != stat.st_size:
        item.size = stat.st_size
        item.save(update_fields=["size", "updated_at"])
    location, created_location = MediaLocation.objects.update_or_create(
        item=item,
        path=str(path),
        defaults={
            "rel_path": str(path.relative_to(folder)),
            "mtime": stat.st_mtime,
        },
    )
    return item, created_item, location, created_location


def rename_location_to_j_id(item, location):
    """Name a freshly indexed copy J-<id>.<ext> inside its own folder."""
    path = Path(location.path)
    if not path.exists():
        return location
    if path.stem == f"J-{item.id}":
        return location
    target = unique_destination(path.parent, f"J-{item.id}{path.suffix}")
    path.rename(target)
    parent = Path(location.rel_path).parent
    location.path = str(target)
    location.rel_path = (
        str(parent / target.name) if str(parent) != "." else target.name
    )
    location.save(update_fields=["path", "rel_path"])
    return location


def compute_visual_hashes(path):
    """Perceptual hashes for an image file (empty dict for other files)."""
    path = Path(path)
    if path.suffix.lower() not in IMAGE_EXTENSIONS:
        return {}
    try:
        with Image.open(path) as image:
            converted = image.convert("RGB")
            return {
                "ahash": str(imagehash.average_hash(converted, hash_size=8)),
                "dhash": str(imagehash.dhash(converted, hash_size=8)),
                "phash": str(imagehash.phash(converted, hash_size=8)),
                "whash": str(imagehash.whash(converted, hash_size=8)),
            }
    except Exception:  # noqa: BLE001 - hashing must never break indexing
        logger.exception("Could not compute visual hashes for %s", path)
        return {}


def ensure_visual_hashes(item):
    """Fill in missing perceptual hashes for a media item."""
    if all(getattr(item, field) for field in HASH_FIELDS):
        return item
    location = item.locations.first()
    if location is None:
        return item
    hashes = compute_visual_hashes(location.path)
    if not hashes:
        return item
    for field, value in hashes.items():
        setattr(item, field, value)
    item.save(update_fields=[*hashes.keys(), "updated_at"])
    return item


def parse_tags(raw):
    """Normalize a comma-separated string or JSON list into a list of tags."""
    if raw is None:
        return None
    if isinstance(raw, (list, tuple)):
        values = raw
    else:
        text = str(raw).strip()
        if not text:
            return None
        if text.startswith("["):
            try:
                values = json.loads(text)
            except json.JSONDecodeError:
                values = text.split(",")
        else:
            values = text.split(",")
    cleaned = []
    for value in values:
        tag = str(value).strip()
        if tag and tag not in cleaned:
            cleaned.append(tag[:100])
    return cleaned


def unique_destination(folder, filename):
    """Return a non-existing path inside folder for the given filename."""
    folder = Path(folder)
    folder.mkdir(parents=True, exist_ok=True)
    filename = get_valid_filename(Path(filename).name) or "upload"
    candidate = folder / filename
    stem, suffix = candidate.stem, candidate.suffix
    counter = 1
    while candidate.exists():
        candidate = folder / f"{stem}-{counter}{suffix}"
        counter += 1
    return candidate


class RangeFileWrapper:
    """Iterate over a limited byte range of an open file."""

    def __init__(self, file, length, chunk_size=64 * 1024):
        self.file = file
        self.remaining = length
        self.chunk_size = chunk_size

    def __iter__(self):
        return self

    def __next__(self):
        if self.remaining <= 0:
            raise StopIteration
        data = self.file.read(min(self.chunk_size, self.remaining))
        if not data:
            raise StopIteration
        self.remaining -= len(data)
        return data

    def close(self):
        self.file.close()


def serve_file(request, path, download=False):
    """Serve a file with HTTP range support (needed for video seeking)."""
    path = Path(path)
    if not path.is_file():
        raise Http404
    size = path.stat().st_size
    content_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
    range_header = request.headers.get("Range", "").strip()
    if range_header:
        match = RANGE_RE.match(range_header)
        if match:
            start_raw, end_raw = match.groups()
            if start_raw == "" and end_raw:
                length = min(int(end_raw), size)
                start, end = size - length, size - 1
            else:
                start = int(start_raw or 0)
                end = min(int(end_raw) if end_raw else size - 1, size - 1)
            if start >= size or start > end:
                response = HttpResponse(status=416)
                response["Content-Range"] = f"bytes */{size}"
                return response
            length = end - start + 1
            handle = open(path, "rb")
            handle.seek(start)
            response = FileResponse(
                RangeFileWrapper(handle, length),
                status=206,
                content_type=content_type,
            )
            response["Content-Length"] = str(length)
            response["Content-Range"] = f"bytes {start}-{end}/{size}"
            response["Accept-Ranges"] = "bytes"
            return response
    response = FileResponse(
        open(path, "rb"),
        content_type=content_type,
        as_attachment=download,
        filename=path.name,
    )
    response["Accept-Ranges"] = "bytes"
    return response


class DownloadCancelled(Exception):
    """Raised when a streamed download is cancelled by the user."""


def download_file(
    url, destination, progress_callback=None, should_cancel=None
):
    """Stream a remote file into destination (used by Download to Library)."""
    import requests

    headers = {"User-Agent": settings.USER_AGENT}
    with requests.get(url, headers=headers, stream=True, timeout=120) as response:
        response.raise_for_status()
        total = int(response.headers.get("content-length") or 0)
        downloaded = 0
        with open(destination, "wb") as handle:
            for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
                if should_cancel is not None and should_cancel():
                    raise DownloadCancelled("download cancelled")
                if chunk:
                    handle.write(chunk)
                    downloaded += len(chunk)
                    if progress_callback is not None:
                        progress_callback(downloaded, total)
        if progress_callback is not None:
            progress_callback(downloaded, total or downloaded)


E621_DESCRIPTION_LIMIT = 20000


def trim_e621_post(post):
    """Keep a compact, size-bounded copy of an e621 post payload."""
    if not isinstance(post, dict):
        return None
    file_data = post.get("file") or {}
    relationships = post.get("relationships") or {}
    score = post.get("score") or {}
    categories = post.get("tags") or {}
    return {
        "id": post.get("id"),
        "created_at": post.get("created_at"),
        "rating": post.get("rating"),
        "tags": {
            str(category): [str(tag)[:200] for tag in tags][:500]
            for category, tags in categories.items()
            if isinstance(tags, list)
        },
        "score": {
            "up": score.get("up"),
            "down": score.get("down"),
            "total": score.get("total"),
        },
        "fav_count": post.get("fav_count"),
        "comment_count": post.get("comment_count"),
        "sources": [
            str(source)[:500]
            for source in (post.get("sources") or [])
            if source
        ][:20],
        "description": str(post.get("description") or "")[:E621_DESCRIPTION_LIMIT],
        "pools": [
            int(pool)
            for pool in (post.get("pools") or [])
            if str(pool).isdigit()
        ][:50],
        "relationships": {
            "parent_id": relationships.get("parent_id"),
            "has_children": relationships.get("has_children"),
        },
        "file": {
            "md5": file_data.get("md5"),
            "ext": file_data.get("ext"),
            "size": file_data.get("size"),
            "width": file_data.get("width"),
            "height": file_data.get("height"),
            "url": file_data.get("url"),
        },
        "uploader_name": post.get("uploader_name"),
    }


def sanitize_iqdb_results(results):
    """Keep a compact, size-bounded copy of IQDB candidates."""
    cleaned = []
    for result in results[:10]:
        if not isinstance(result, dict):
            continue
        score = result.get("score")
        post_id = result.get("post_id")
        tags_preview = result.get("tags_preview")
        entry = {
            "post_id": post_id if isinstance(post_id, int) else None,
            "score": float(score) if isinstance(score, (int, float)) else None,
            "preview_url": str(result.get("preview_url") or "")[:500] or None,
            "rating": str(result.get("rating") or "")[:1] or None,
            "md5": str(result.get("md5") or "")[:32] or None,
            "score_total": (
                result.get("score_total")
                if isinstance(result.get("score_total"), int)
                else None
            ),
            "fav_count": (
                result.get("fav_count")
                if isinstance(result.get("fav_count"), int)
                else None
            ),
            "width": (
                result.get("width")
                if isinstance(result.get("width"), int)
                else None
            ),
            "height": (
                result.get("height")
                if isinstance(result.get("height"), int)
                else None
            ),
            "tags_preview": (
                [str(tag)[:100] for tag in tags_preview][:8]
                if isinstance(tags_preview, list)
                else []
            ),
        }
        if entry["post_id"] or entry["preview_url"]:
            cleaned.append(entry)
    return cleaned


def generate_video_thumbnail(md5, path):
    """Extract a JPEG thumbnail from a video, cached under MEDIA_ROOT/thumbs."""
    if not shutil.which("ffmpeg"):
        return None
    thumbs_dir = Path(settings.MEDIA_ROOT) / "thumbs"
    thumbs_dir.mkdir(parents=True, exist_ok=True)
    target = thumbs_dir / f"{md5}.jpg"
    if target.exists() and target.stat().st_mtime >= os.path.getmtime(path):
        return target
    command = [
        "ffmpeg",
        "-y",
        "-ss",
        "0.5",
        "-i",
        str(path),
        "-frames:v",
        "1",
        "-vf",
        "scale=480:-2",
        "-loglevel",
        "error",
        str(target),
    ]
    try:
        subprocess.run(command, check=True, capture_output=True, timeout=60)
    except (subprocess.SubprocessError, OSError):
        return None
    return target if target.exists() else None