J621/backend/apps/library/services.py

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

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

from .models import MediaItem, MediaLocation

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, 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, created_location


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


def download_file(url, destination):
    """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()
        with open(destination, "wb") as handle:
            for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
                if chunk:
                    handle.write(chunk)


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