yunke/yunke-server/app/scraper/media.py

"""Media downloader (M2, Task 2.6).

Fetches WeChat / OSS media blobs referenced by message records into the local
:class:`app.storage.object_store.ObjectStore` and records them in the
:class:`app.db.models.Media` table, deduplicated by the source ``md5``.

Phase-0 contract (see ``docs/superpowers/specs/phase0-findings.md`` §2/§3):

* WeChat type ``2`` (image / emoji, ``.webp``) and ``4`` (video, ``.mp4``) carry
  a DIRECT http(s) URL on ``file`` / ``fileTh`` / ``fileHd`` plus ``md5`` ->
  downloadable, deduped by ``md5``.
* WeChat type ``3`` (legacy image) carries an OPAQUE id on ``file`` (not a URL)
  -> cannot be resolved yet; :func:`download_message_media` returns ``None``.
* WeChat type ``10`` (system / link / card / invite) carries a
  ``support.weixin.qq.com`` URL that is a documentation link, NOT a media blob
  -> explicitly skipped before any network attempt.

The downloader is URL-driven and generic: any http(s) media URL works (the
recording URLs on ``yunke-pcfile`` are the same shape), so call recordings and
voice-call recordings can be fetched by the same path once their owning records
are wired in. That wiring is OUT of scope here (this module only consumes
normalized message dicts); see ``app.scraper.wechat`` for those.

Resilient by design: a download failure is logged and the function returns
``None`` rather than raising. A retry queue is a LATER milestone.
"""

from __future__ import annotations

import logging
import mimetypes
from pathlib import PurePosixPath
from typing import Any
from urllib.parse import urlparse

import httpx
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models import Media, Message, MomentsPost, PhoneCall, VoiceClip
from app.session.yunke_client import YunkeClient
from app.storage.object_store import ObjectStore
from app.storage.repos import upsert_media

__all__ = [
    "download_call_recordings",
    "download_message_images",
    "download_message_media",
    "download_moments_images",
    "download_url",
    "download_voice_clip_audio",
    "guess_mime_type",
    "media_type_from_message",
]

logger = logging.getLogger(__name__)

# Normalized message ``type`` label -> :attr:`Media.media_type` value. Only the
# Phase-0-confirmed image / video labels map today; voice / file codes
# (``6`` / ``9`` / ...) remain ``"unknown"`` at the scraper layer and yield
# ``None`` here until they are confirmed (phase0-findings.md §2).
_MEDIA_TYPE_BY_LABEL: dict[str, str] = {
    "image": "image",
    "image_legacy": "image",
    "video": "video",
}

# Extension -> MIME fallback map for the OSS / call-recording media types we
# observe. Used when the stdlib ``mimetypes`` database does not recognize the
# extension (e.g. ``.webp`` / ``.amr`` on some platforms) so content-type
# derivation does not depend on the system MIME database.
_EXT_MIME: dict[str, str] = {
    ".webp": "image/webp",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
    ".gif": "image/gif",
    ".bmp": "image/bmp",
    ".mp4": "video/mp4",
    ".mp3": "audio/mpeg",
    ".m4a": "audio/mp4",
    ".wav": "audio/wav",
    ".amr": "audio/amr",
}


def media_type_from_message(msg: dict[str, Any]) -> str | None:
    """Map a normalized message dict to a :attr:`Media.media_type` label.

    Returns ``"image"`` for ``image`` / ``image_legacy`` messages and ``"video"``
    for ``video`` messages; ``None`` for everything else (text / system /
    unknown / not-yet-confirmed voice-or-file codes).
    """
    return _MEDIA_TYPE_BY_LABEL.get(msg.get("type"))  # type: ignore[arg-type]


def _is_http_url(value: Any) -> bool:
    """True iff ``value`` is an ``http://`` / ``https://`` string."""
    return isinstance(value, str) and value.startswith(("http://", "https://"))


def _pick_media_url(msg: dict[str, Any]) -> str | None:
    """Pick the best downloadable http(s) URL from a message dict.

    Preference order: ``file_url`` (original) -> ``file_hd_url`` (high-def) ->
    ``file_th_url`` (thumbnail). Returns ``None`` when no candidate is an
    http(s) URL (covers type-``3`` opaque ids and text messages).
    """
    for key in ("file_url", "file_hd_url", "file_th_url"):
        url = msg.get(key)
        if _is_http_url(url):
            return url  # type: ignore[return-value]
    return None


def _ext_from_url(url: str) -> str:
    """Extract the lowercase extension (with leading ``.``) from a URL path.

    Query / fragment are stripped first. Returns ``""`` when the path has no
    suffix (e.g. an OSS callback URL without an extension).
    """
    suffix = PurePosixPath(urlparse(url).path).suffix
    return suffix.lower()


def guess_mime_type(ext: str) -> str | None:
    """Guess a MIME type from a file extension.

    The built-in :data:`_EXT_MIME` map is consulted first (covers ``.webp`` /
    ``.mp4`` / ``.amr`` reliably regardless of platform), then the stdlib
    ``mimetypes`` database. Returns ``None`` when nothing matches.
    """
    if not ext:
        return None
    key = ext.lower()
    if key in _EXT_MIME:
        return _EXT_MIME[key]
    guessed, _ = mimetypes.guess_type(f"x{key}")
    return guessed


async def download_url(client: YunkeClient, url: str) -> bytes:
    """Fetch the raw bytes at ``url`` via the client's underlying httpx session.

    The :class:`YunkeClient` httpx client carries the session cookies (harmless
    for the public OSS media URLs, useful for any same-origin authenticated
    URLs) and is the transport ``respx`` mocks in the test suite, so media
    downloads go through the same mocked transport as every other scraper
    request.

    Redirects are NOT followed (SSRF protection): the URL is the trusted OSS /
    recording endpoint resolved at scrape time, and following a 3xx would allow
    a malicious or compromised host to redirect us to internal addresses such
    as ``http://169.254.169.254/`` (cloud metadata) or other RFC1918 hosts.
    OSS media URLs return the blob directly with 200; a 3xx here is treated as
    a download failure and surfaces through the caller's existing
    ``except (httpx.HTTPError, OSError)`` clause.
    """
    resp = await client.client.get(url, follow_redirects=False)
    resp.raise_for_status()
    return resp.content


async def download_message_media(
    client: YunkeClient,
    msg: dict[str, Any],
    store: ObjectStore,
    session: AsyncSession,
) -> str | None:
    """Download one message's media blob and record it in the ``Media`` table.

    Returns the dedup hash (the message ``md5`` when present, else the sha256 of
    the downloaded bytes), or ``None`` when:

    * the message is a type-``10`` system / link message (non-media; skipped
      before any network attempt), or
    * the message carries no downloadable http(s) URL (text, type-``3``
      opaque-id images, or any message whose ``file``/``fileHd``/``fileTh`` are
      absent or non-URL), or
    * the download fails (the error is logged; no exception escapes).

    Dedup: when the message carries an ``md5`` and a :class:`Media` row for that
    hash already exists with a local file, the download is SKIPPED and the hash
    is returned without touching the network.

    The blob is written via :meth:`ObjectStore.put_by_hash` (sha256
    content-addressed on disk), so identical bytes are stored once regardless of
    how many messages reference them.
    """
    # Type-10 system / link messages carry a support URL that is NOT media
    # (phase0-findings.md §2): skip them explicitly before URL selection so the
    # downloader never fetches a documentation page as if it were a media blob.
    if msg.get("type") == "system":
        return None

    url = _pick_media_url(msg)
    if url is None:
        return None

    md5 = msg.get("md5")
    if md5:
        existing = (
            await session.execute(
                sa.select(Media.local_path).where(Media.media_hash == md5)
            )
        ).scalar_one_or_none()
        if existing:
            # Already downloaded (local file present); short-circuit the network.
            return md5  # type: ignore[return-value]

    try:
        data = await download_url(client, url)
    except (httpx.HTTPError, OSError) as exc:
        logger.warning("media download failed for %s: %s", url, exc)
        return None

    ext = _ext_from_url(url)
    sha256_hex, path = await store.put_by_hash(data, ext=ext)
    media_hash: str = md5 if md5 else sha256_hex

    await upsert_media(
        session,
        media_hash=media_hash,
        channel=msg.get("channel") or "wechat",
        media_type=media_type_from_message(msg),
        original_url=msg.get("file_url"),
        original_th_url=msg.get("file_th_url"),
        original_hd_url=msg.get("file_hd_url"),
        local_path=str(path),
        size_bytes=len(data),
        mime_type=guess_mime_type(ext),
        raw_json=msg.get("raw_json"),
    )
    return media_hash


async def download_call_recordings(
    client: YunkeClient,
    store: ObjectStore,
    session: AsyncSession,
    *,
    batch_size: int = 50,
) -> tuple[int, int]:
    """Download call recordings (``phone_call.record_file`` OSS URLs) into the
    local object store and link them via ``phone_call.record_media_hash``.

    Iterates ``phone_call`` rows that carry an http(s) ``record_file`` but no
    ``record_media_hash`` yet. For each: fetch bytes (:func:`download_url`),
    store content-addressed (:meth:`ObjectStore.put_by_hash`, sha256 dedup),
    upsert a ``Media`` row (``media_type='audio'``), and write
    ``record_media_hash`` back onto the call.

    Returns ``(attempted, localized)``. Resilient: a per-row download/store
    failure is logged and counted as attempted-but-not-localized; no exception
    escapes. Idempotent within a run: a per-run ``tried`` set excludes already-
    attempted call_ids from later batches so a persistently-failing row is NOT
    re-fetched infinitely (it stays ``record_media_hash IS NULL`` and is retried
    on the NEXT run, which is desirable). Rows already carrying
    ``record_media_hash`` are skipped; identical bytes are stored once.
    """
    attempted = 0
    localized = 0
    tried: set[str] = set()
    while True:
        stmt = (
            sa.select(PhoneCall.call_id, PhoneCall.record_file, PhoneCall.channel)
            .where(
                PhoneCall.record_file.is_not(None),
                PhoneCall.record_file.like("http%"),
                PhoneCall.record_media_hash.is_(None),
            )
            .limit(batch_size)
        )
        if tried:
            stmt = stmt.where(PhoneCall.call_id.not_in(tried))
        rows = (await session.execute(stmt)).all()
        if not rows:
            break
        for call_id, record_file, channel in rows:
            attempted += 1
            tried.add(call_id)
            try:
                data = await download_url(client, record_file)
            except (httpx.HTTPError, OSError) as exc:
                logger.warning("recording download failed for %s: %s", call_id, exc)
                continue
            ext = _ext_from_url(record_file)
            sha256_hex, path = await store.put_by_hash(data, ext=ext)
            await upsert_media(
                session,
                media_hash=sha256_hex,
                channel=channel or "wechat",
                media_type="audio",
                original_url=record_file,
                local_path=str(path),
                size_bytes=len(data),
                mime_type=guess_mime_type(ext),
            )
            await session.execute(
                sa.update(PhoneCall)
                .where(PhoneCall.call_id == call_id)
                .values(record_media_hash=sha256_hex)
            )
            localized += 1
        await session.commit()
    return attempted, localized


async def download_voice_clip_audio(
    client: YunkeClient,
    store: ObjectStore,
    session: AsyncSession,
    *,
    batch_size: int = 50,
) -> tuple[int, int]:
    """Download voice_clip audio (``voice_clip.oss_file_name`` OSS URLs) into
    the local object store and link them via ``voice_clip.media_hash``.

    Mirrors :func:`download_call_recordings` exactly: iterates ``voice_clip``
    rows that carry an http(s) ``oss_file_name`` but no ``media_hash`` yet. For
    each: fetch bytes (:func:`download_url`), store content-addressed
    (:meth:`ObjectStore.put_by_hash`, sha256 dedup), upsert a ``Media`` row
    (``media_type='audio'``), and write ``media_hash`` back onto the clip.

    Returns ``(attempted, localized)``. Resilient: a per-row download/store
    failure is logged and counted as attempted-but-not-localized; no exception
    escapes. Idempotent within a run: a per-run ``tried`` set excludes already-
    attempted voice_clip_ids from later batches so a persistently-failing row
    is NOT re-fetched infinitely (it stays ``media_hash IS NULL`` and is
    retried on the NEXT run). Rows already carrying ``media_hash`` are skipped;
    identical bytes are stored once.
    """
    attempted = 0
    localized = 0
    tried: set[int] = set()
    while True:
        stmt = (
            sa.select(
                VoiceClip.voice_clip_id,
                VoiceClip.oss_file_name,
                VoiceClip.channel,
            )
            .where(
                VoiceClip.oss_file_name.is_not(None),
                VoiceClip.oss_file_name.like("http%"),
                VoiceClip.media_hash.is_(None),
            )
            .limit(batch_size)
        )
        if tried:
            stmt = stmt.where(VoiceClip.voice_clip_id.not_in(tried))
        rows = (await session.execute(stmt)).all()
        if not rows:
            break
        for voice_clip_id, oss_file_name, channel in rows:
            attempted += 1
            tried.add(voice_clip_id)
            try:
                data = await download_url(client, oss_file_name)
            except (httpx.HTTPError, OSError) as exc:
                logger.warning(
                    "voice_clip audio download failed for %s: %s",
                    voice_clip_id,
                    exc,
                )
                continue
            ext = _ext_from_url(oss_file_name)
            sha256_hex, path = await store.put_by_hash(data, ext=ext)
            await upsert_media(
                session,
                media_hash=sha256_hex,
                channel=channel or "wechat",
                media_type="audio",
                original_url=oss_file_name,
                local_path=str(path),
                size_bytes=len(data),
                mime_type=guess_mime_type(ext),
            )
            await session.execute(
                sa.update(VoiceClip)
                .where(VoiceClip.voice_clip_id == voice_clip_id)
                .values(media_hash=sha256_hex)
            )
            localized += 1
        await session.commit()
    return attempted, localized


def _pick_message_image_url(
    *,
    type_: str | None,
    file_url: str | None,
    file_th_url: str | None,
    file_hd_url: str | None,
) -> str | None:
    """Pick the best downloadable http(s) URL for a message media row.

    For ``image`` / ``image_legacy`` messages, prefer the thumbnail
    (``file_th_url``): thumbnails are small, fast to fetch, and exactly what the
    chat timeline renders, so localizing the thumbnail is sufficient for display
    and easy on disk + bandwidth. Falls back to ``file_url`` then ``file_hd_url``
    when no http(s) thumbnail is available. ``image_legacy`` (smoke.db 26 rows
    with valid ``wx.oss.yunkecn.com`` URLs) is treated exactly like ``image``
    here so the legacy WeChat type-3-via-OSS rows are localized.

    For every other media type (``video`` / ``file``), prefer ``file_url`` (the
    original), then ``file_hd_url``, then ``file_th_url``.

    Returns ``None`` when no candidate is an http(s) URL (e.g. type-3 opaque ids
    or relative filenames). Only ``http://`` / ``https://`` URLs are downloadable.
    """
    if type_ in ("image", "image_legacy"):
        order = (file_th_url, file_url, file_hd_url)
    else:
        order = (file_url, file_hd_url, file_th_url)
    for url in order:
        if _is_http_url(url):
            return url  # type: ignore[return-value]
    return None


async def download_message_images(
    client: YunkeClient,
    store: ObjectStore,
    session: AsyncSession,
    *,
    batch_size: int = 50,
) -> tuple[int, int]:
    """Download chat-message media (image / video / file) OSS URLs into the
    local object store and link them via ``message.media_hash``.

    Iterates ``message`` rows that carry an http(s) URL on any of
    ``file_url`` / ``file_th_url`` / ``file_hd_url``, have ``type`` in
    ('image', 'image_legacy', 'video', 'file'), are on the wechat / whatsapp
    channel, and do NOT yet carry a ``media_hash``. For each: pick the best
    available URL (image: thumbnail-first; other: original-first — see
    :func:`_pick_message_image_url`), fetch bytes (:func:`download_url`),
    store content-addressed (:meth:`ObjectStore.put_by_hash`, sha256 dedup),
    upsert a ``Media`` row, and write ``media_hash`` back onto the message.
    ``image_legacy`` rows are stored as ``media_type='image'`` (normalized via
    :func:`media_type_from_message`) so the legacy WeChat type-3-via-OSS rows
    (smoke.db 26 rows with valid ``wx.oss.yunkecn.com`` URLs) are covered.

    Returns ``(attempted, localized)``. Resilient: a per-row download/store
    failure is logged and counted as attempted-but-not-localized; no exception
    escapes. Idempotent within a run: a per-run ``tried`` set excludes already-
    attempted msg_ids from later batches so a persistently-failing row is NOT
    re-fetched infinitely (it stays ``media_hash IS NULL`` and is retried on
    the NEXT run). Rows already carrying ``media_hash`` are skipped; identical
    bytes are stored once.

    Mirrors :func:`download_call_recordings` / :func:`download_voice_clip_audio`
    exactly. The message ``type`` label filter restricts to media-bearing rows
    so text / system / unknown messages are never selected.
    """
    attempted = 0
    localized = 0
    tried: set[str] = set()
    while True:
        stmt = (
            sa.select(
                Message.msg_id,
                Message.type,
                Message.channel,
                Message.file_url,
                Message.file_th_url,
                Message.file_hd_url,
            )
            .where(Message.media_hash.is_(None))
            .where(Message.channel.in_(("wechat", "whatsapp")))
            .where(Message.type.in_(("image", "image_legacy", "video", "file")))
            .where(
                sa.or_(
                    Message.file_url.like("http%"),
                    Message.file_th_url.like("http%"),
                    Message.file_hd_url.like("http%"),
                )
            )
            .limit(batch_size)
        )
        if tried:
            # 分片进 NOT IN —— 单次塞全部 tried(3万+)会超 SQLite 变量上限
            # (默认 32766/999),2026-08-20 实测 42k 后崩;每次最多 500 个。
            tried_list = list(tried)
            conds = []
            for i in range(0, len(tried_list), 500):
                conds.append(Message.msg_id.not_in(tried_list[i:i + 500]))
            stmt = stmt.where(sa.and_(*conds))
        rows = (await session.execute(stmt)).all()
        if not rows:
            break
        for msg_id, type_, channel, file_url, file_th_url, file_hd_url in rows:
            attempted += 1
            tried.add(msg_id)
            url = _pick_message_image_url(
                type_=type_,
                file_url=file_url,
                file_th_url=file_th_url,
                file_hd_url=file_hd_url,
            )
            if url is None:
                # WHERE clause matched at least one http(s) URL, but the
                # preferred slot for this type isn't http(s) — defensive, in
                # practice the OR above guarantees _pick returns a value.
                continue
            try:
                data = await download_url(client, url)
            except (httpx.HTTPError, OSError) as exc:
                logger.warning(
                    "message media download failed for %s: %s", msg_id, exc
                )
                continue
            ext = _ext_from_url(url)
            sha256_hex, path = await store.put_by_hash(data, ext=ext)
            # Normalize ``image_legacy`` to ``image`` on the stored Media row so
            # the legacy WeChat type-3-via-OSS rows land in the same bucket as
            # type-2 images (matches ``media_type_from_message``).
            media_type = "image" if type_ == "image_legacy" else type_
            await upsert_media(
                session,
                media_hash=sha256_hex,
                channel=channel or "wechat",
                media_type=media_type,
                original_url=file_url,
                original_th_url=file_th_url,
                original_hd_url=file_hd_url,
                local_path=str(path),
                size_bytes=len(data),
                mime_type=guess_mime_type(ext),
            )
            await session.execute(
                sa.update(Message)
                .where(Message.msg_id == msg_id)
                .values(media_hash=sha256_hex)
            )
            localized += 1
        await session.commit()
    return attempted, localized


async def download_moments_images(
    client: YunkeClient,
    store: ObjectStore,
    session: AsyncSession,
    *,
    batch_size: int = 50,
) -> tuple[int, int]:
    """Download 员工朋友圈 images (``moments_post.image_urls`` OSS 直链) into
    the local object store, one ``Media`` row per blob.

    朋友圈与消息图片的关键差异(2026-08-25 接入时确认):``image_urls`` 是
    JSON **数组**(一帖多图),且 :class:`MomentsPost` 没有 ``media_hash``
    回链列 —— 已下载标记改用 ``Media.original_url``:URL 已有对应 Media 行
    即视为已本地化并跳过(不重下载)。无回链也就无法"标记"帖子行,故候选
    集一次物化(全表展开 + 去重),再按 ``original_url`` 剔除已下载的,分批
    下载(批间 commit,部分进度可保留)。

    Returns ``(attempted, localized)``. Resilient: a per-URL download failure is
    logged and counted as attempted-but-not-localized; no exception escapes.
    Failed URLs keep no Media row → retried on the NEXT run. Identical bytes are
    stored once (:meth:`ObjectStore.put_by_hash`, sha256 dedup).

    Known edge: two DISTINCT URLs hosting byte-identical content share one
    ``media_hash`` (content-addressed PK), and the upsert refreshes
    ``original_url`` to the latest — so the loser URL is re-fetched once per run
    (bounded, never grows; identical bytes at two OSS paths is rare).
    """
    rows = (
        await session.execute(
            sa.select(MomentsPost.id, MomentsPost.image_urls).where(
                MomentsPost.image_urls.is_not(None)
            )
        )
    ).all()

    # 一帖多图 → 展开去重(同图多帖只下载一次)
    urls: list[str] = []
    seen: set[str] = set()
    for _post_id, image_urls in rows:
        for url in image_urls or []:
            if _is_http_url(url) and url not in seen:
                seen.add(url)
                urls.append(url)
    if not urls:
        return 0, 0

    # 剔除已本地化的 URL。分片 IN —— 单次塞全部 URL 会超 SQLite 变量上限
    # (同 download_message_images 的 tried 分片惯例,2026-08-20)。
    existing: set[str] = set()
    for i in range(0, len(urls), 500):
        res = await session.execute(
            sa.select(Media.original_url).where(
                Media.original_url.in_(urls[i : i + 500])
            )
        )
        existing.update(u for (u,) in res.fetchall() if u)
    pending = [u for u in urls if u not in existing]

    attempted = 0
    localized = 0
    for i in range(0, len(pending), batch_size):
        for url in pending[i : i + batch_size]:
            attempted += 1
            try:
                data = await download_url(client, url)
            except (httpx.HTTPError, OSError) as exc:
                logger.warning("moments image download failed for %s: %s", url, exc)
                continue
            ext = _ext_from_url(url)
            sha256_hex, path = await store.put_by_hash(data, ext=ext)
            await upsert_media(
                session,
                media_hash=sha256_hex,
                channel="wechat",  # 朋友圈是微信独有实体(无 WA 朋友圈)
                media_type="image",
                original_url=url,
                local_path=str(path),
                size_bytes=len(data),
                mime_type=guess_mime_type(ext),
            )
            localized += 1
        await session.commit()
    return attempted, localized