FabledCurator/backend/app/services/pixiv_downloader.py

"""Native Pixiv media downloader — the Pixiv counterpart to
patreon_downloader / subscribestar_downloader.

Given a normalized Pixiv work and its resolved `MediaItem`s
(pixiv_client.extract_media), download the originals to gallery-dl's on-disk
layout (so pre-cutover gallery-dl downloads are recognized on disk and not
re-fetched), write the post-first sidecars the importer consumes, and report
per-media outcomes.

On-disk layout (matches FC's gallery-dl pixiv config, PLATFORM_DEFAULTS:
base-directory `<images_root>/<artist_slug>/pixiv` + `directory:
["{category}"]` + filename `{id}_{title[:50]}_{num:>02}.{extension}`):

    <images_root>/<artist_slug>/pixiv/pixiv/<id>_<title50>_<NN>.<ext>

— note the intentional DOUBLE `pixiv` segment: gallery-dl appended
`{category}` under a base-directory that already ended in the platform name,
and tier-2 disk-skip parity requires reproducing that exactly. The layout is
FLAT (no per-post directory), so the post-first record is `_post_<id>.json`
in the same directory (the id suffix prevents the collisions a bare
`_post.json` would have here; phase 3 receives explicit post_record_paths, so
the name is a convention, not a discovery key).

Simpler than Patreon (no Mux/yt-dlp video branch) — the one special file is
the ugoira frame zip, downloaded as-is; FC's archive-containment import
extracts the frames, and the frame DELAYS ride the post record (the zip
carries none — a future ugoira→video conversion needs them).

PURE: no DB; the seen-skip is an injected predicate. FC runs on a plain-HTTP
homelab; nothing here uses a secure-context Web API.
"""

from __future__ import annotations

import json
import logging
import time
from collections.abc import Callable
from pathlib import Path

import requests

from .native_ingest_common import (
    BaseNativeDownloader,
    MediaOutcome,
    PostRecordOutcome,
    sanitize_segment,
)
from .pixiv_client import PIXIV_APP_HEADERS, rating_label

log = logging.getLogger(__name__)

# Enrichment keys copied verbatim from the app-API work dict into the post
# record (they're already JSON scalars/objects). Everything lands in
# Post.raw_metadata via the importer, so the archive keeps pixiv's stats and
# structure without a schema change.
_WORK_PASSTHROUGH_KEYS = (
    "type",
    "page_count",
    "width",
    "height",
    "total_view",
    "total_bookmarks",
    "total_comments",
    "is_bookmarked",
    "illust_ai_type",
    "series",
)


class PixivDownloader(BaseNativeDownloader):
    """Download resolved Pixiv media to gallery-dl's on-disk layout.
    Subclasses BaseNativeDownloader for the shared streaming GET
    (transient-retry + Range-resume) and validation/quarantine. PURE: no DB."""

    def __init__(
        self,
        images_root: Path,
        cookies_path: str | None = None,
        *,
        validate: bool = True,
        rate_limit: float = 0.0,
        session: requests.Session | None = None,
    ):
        super().__init__(
            images_root, cookies_path, platform="pixiv",
            validate=validate, rate_limit=rate_limit, session=session,
        )
        if session is None:
            # i.pximg.net 403s any GET without the app Referer; mirror the
            # client's full app-header profile (gallery-dl serves media off
            # the same session it drives the API with). An injected session
            # (tests) owns its own headers.
            self.session.headers.update(PIXIV_APP_HEADERS)

    # -- public ------------------------------------------------------------

    def download_post(
        self,
        post: dict,
        media_items: list,
        artist_slug: str,
        *,
        is_seen: Callable[[object], bool] = lambda m: False,
        should_stop: Callable[[], bool] = lambda: False,
        recapture: bool = False,
    ) -> list[MediaOutcome]:
        """Download every media item of one work; return per-item outcomes.
        Mirrors SubscribeStarDownloader.download_post (two-tier skip, mid-post
        time-box, recapture surfacing)."""
        flat_dir = self._flat_dir(artist_slug)
        outcomes: list[MediaOutcome] = []
        for media in media_items:
            if should_stop():
                break
            try:
                outcomes.append(
                    self._download_one(
                        post, media, flat_dir, artist_slug, is_seen,
                        recapture=recapture,
                    )
                )
            except Exception as exc:  # resilient: isolate one item's failure
                log.warning(
                    "Pixiv media failed (work %s, %s): %s",
                    post.get("id"), getattr(media, "media_id", "?"), exc,
                )
                outcomes.append(
                    MediaOutcome(media=media, status="error", path=None, error=str(exc))
                )
        return outcomes

    def _flat_dir(self, artist_slug: str) -> Path:
        # Double platform segment — gallery-dl layout parity (module docstring).
        return self.images_root / artist_slug / "pixiv" / "pixiv"

    # -- per-item ----------------------------------------------------------

    def _download_one(
        self,
        post: dict,
        media,
        flat_dir: Path,
        artist_slug: str,
        is_seen: Callable[[object], bool],
        *,
        recapture: bool = False,
    ) -> MediaOutcome:
        seen = is_seen(media)
        if seen and not recapture:
            return MediaOutcome(media=media, status="skipped_seen", path=None, error=None)

        # The client's filename already carries the {id}_{title50}_{NN} shape;
        # only path-restrict sanitization happens here.
        media_path = flat_dir / sanitize_segment(media.filename)

        if media_path.exists():  # tier-2: already on disk
            return MediaOutcome(
                media=media, status="skipped_disk", path=media_path, error=None
            )
        # recapture: a seen item not on disk is NOT re-downloaded (recovery's job).
        if seen:
            return MediaOutcome(media=media, status="skipped_seen", path=None, error=None)

        flat_dir.mkdir(parents=True, exist_ok=True)
        if self._rate_limit > 0:
            time.sleep(self._rate_limit)

        out_path = self._fetch_get(media.url, media_path)
        reason, quarantine_dest = self._validate_path(out_path, artist_slug, media.url)
        if reason is not None:
            return MediaOutcome(
                media=media, status="quarantined", path=quarantine_dest, error=reason,
            )
        self._write_minimal_sidecar(post, out_path, source_url=media.url)
        return MediaOutcome(media=media, status="downloaded", path=out_path, error=None)

    # -- post record ---------------------------------------------------------

    def write_post_record(self, post: dict, artist_slug: str) -> PostRecordOutcome:
        """Write the post-first `_post_<id>.json` — the sole writer of the post
        body/metadata on the native path. Beyond the standard body fields, the
        record carries pixiv's own structure (tags + EN translations, rating,
        series, view/bookmark counts, AI flag, dimensions, author, ugoira frame
        delays) so the archive keeps what the platform knows about the work."""
        attrs = post.get("attributes") or {}
        work = post.get("_work") or {}
        title = attrs.get("title") if isinstance(attrs.get("title"), str) else None
        post_type = attrs.get("post_type") if isinstance(attrs.get("post_type"), str) else None
        pid = str(post.get("id") or "")
        if not pid:
            return PostRecordOutcome(
                path=None, post_type=post_type, title=title, body_chars=0,
            )

        content = attrs.get("content")
        content = content if isinstance(content, str) else ""
        data: dict = {
            "category": "pixiv",
            "id": pid,
            "title": title or "",
            "content": content,
            "published_at": attrs.get("published_at"),
            # The post permalink is synthesized by platforms/pixiv.py
            # derive_post_url from `id` at parse time — no url key here.
            "rating": rating_label(work.get("x_restrict")),
        }
        for key in _WORK_PASSTHROUGH_KEYS:
            if key in work:
                data[key] = work[key]
        tags = work.get("tags")
        if isinstance(tags, list):
            data["tags"] = [
                {
                    "name": t.get("name"),
                    "translated_name": t.get("translated_name"),
                }
                for t in tags
                if isinstance(t, dict)
            ]
        user = work.get("user")
        if isinstance(user, dict):
            data["user"] = {
                "id": user.get("id"),
                "account": user.get("account"),
                "name": user.get("name"),
            }
        # Memoized by extract_media's ugoira branch; the zip has no timings.
        frames = work.get("_ugoira_frames")
        if frames:
            data["ugoira_frames"] = frames

        flat_dir = self._flat_dir(artist_slug)
        flat_dir.mkdir(parents=True, exist_ok=True)
        path = flat_dir / f"_post_{pid}.json"
        path.write_text(json.dumps(data, indent=2, ensure_ascii=False))
        return PostRecordOutcome(
            path=path, post_type=post_type, title=title, body_chars=len(content),
        )