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

"""WeChat scrapers.

Two endpoints today:

* :func:`scrape_wechat_friends` — wraps ``POST /wechatweb/wxFriends/friend_list/es``
  and yields normalized dicts ready to be upserted into
  :class:`app.db.models.Customer`.
* :func:`scrape_wechat_messages` — wraps ``POST /wechatweb/msg/chats`` and
  yields normalized chat-message dicts (see Task 1.6).

Both yield dicts (not ORM). Persistence is a LATER task — Task 1.7. Each dict
carries:

* snake_case keys matching the target ORM columns (see :func:`_normalize`
  for friends and :func:`_normalize_message` for messages),
* ``channel == "wechat"``,
* ``raw_json`` — the original upstream record, verbatim, so nothing captured
  is lost even when a column has no target field.

Friend field names are mapped from the real capture (see Task 1.5 spec). The
mapping is a straight ``camelCase -> snake_case`` translation with a couple of
disambiguations where the source has two fields for one column:

* ``phone``       <- ``phone``
* ``wechat_phone`` <- ``wechatPhone``
* ``dept_name``   <- ``deptName`` (falls back to ``deptRouteName``)
* ``head_image_url`` <- ``friendHeadImageUrl`` (falls back to ``headImg``)
"""

from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any

from app.scraper.base import paginate_by_cursor, paginate_by_page
from app.scraper.parsers import normalize_wechat_type
from app.session.yunke_client import YunkeClient

__all__ = ["scrape_wechat_friends", "scrape_wechat_messages"]

FRIEND_LIST_PATH = "/wechatweb/wxFriends/friend_list/es"
MSG_CHATS_PATH = "/wechatweb/msg/chats"

# Upstream sends form-encoded bodies for ``msg/chats`` (NOT JSON), so we have
# to override the client's default JSON content-type on every request.
_FORM_CONTENT_TYPE = "application/x-www-form-urlencoded"


def _parse_total(resp: Any) -> int | None:
    """Read the declared row count from the friend_list envelope.

    Shape: ``data.basicResult.total`` (an int). Returns ``None`` if the path
    is absent so the paginator falls back to the empty-page stop signal.
    """
    try:
        body = resp.json()
    except (AttributeError, ValueError):
        return None
    data = body.get("data") if isinstance(body, dict) else None
    if not isinstance(data, dict):
        return None
    basic = data.get("basicResult")
    if not isinstance(basic, dict):
        return None
    total = basic.get("total")
    return int(total) if isinstance(total, (int, float)) else None


def _parse_rows(resp: Any) -> list[dict[str, Any]]:
    """Extract the list of friend records from the envelope.

    Shape: ``data.basicResult.data`` (a list of records). Missing/non-list
    values collapse to an empty list so the paginator sees "empty page".
    """
    try:
        body = resp.json()
    except (AttributeError, ValueError):
        return []
    data = body.get("data") if isinstance(body, dict) else None
    if not isinstance(data, dict):
        return []
    basic = data.get("basicResult")
    if not isinstance(basic, dict):
        return []
    rows = basic.get("data")
    return rows if isinstance(rows, list) else []


def _build_body(
    page_num: int,
    page_size: int,
    *,
    department_id: str = "",
    user_id: str = "",
) -> dict[str, Any]:
    """Build the friend_list request body for a given page.

    The real client sends ``departmentId`` / ``userId`` filters; we accept
    them as optional kwargs (default empty) so per-staff / per-department
    scrapes are a drop-in without changing the paginator contract.
    """
    body: dict[str, Any] = {"pageNum": page_num, "pageSize": page_size}
    if department_id:
        body["departmentId"] = department_id
    if user_id:
        body["userId"] = user_id
    return body


def _normalize(record: dict[str, Any]) -> dict[str, Any]:
    """Map one upstream friend record to the ``Customer`` column shape.

    Adds ``channel="wechat"`` and preserves the original record under
    ``raw_json``. ``None``-safe: missing fields become ``None`` rather than
    raising, so partial payloads do not break the whole scrape.
    """
    return {
        "friend_id": record.get("friendWxId"),
        "channel": "wechat",
        "owner_account_id": record.get("wxId"),
        "friend_nickname": record.get("friendWxNickname"),
        "friend_remark": record.get("friendWxRemark"),
        "real_name": record.get("realName"),
        "phone": record.get("phone"),
        "wechat_phone": record.get("wechatPhone"),
        "gender": record.get("gender"),
        "region": record.get("region"),
        "description": record.get("description"),
        "sale_name": record.get("saleName"),
        "user_id": record.get("userId"),
        "dept_route": record.get("deptRoute"),
        # Prefer the canonical deptName; fall back to the route-display name.
        "dept_name": record.get("deptName") or record.get("deptRouteName"),
        # Prefer friendHeadImageUrl; fall back to the legacy headImg field.
        "head_image_url": record.get("friendHeadImageUrl") or record.get("headImg"),
        "last_chat_time": record.get("lastChatTime"),
        "send_num": record.get("sendNum"),
        "receive_num": record.get("receiveNum"),
        "message_count": record.get("messageCount"),
        "delete_flag": record.get("deleteFlag"),
        "is_block": record.get("isBlock"),
        "friend_state": record.get("friendState"),
        "version": record.get("version"),
        "raw_json": record,
    }


async def scrape_wechat_friends(
    client: YunkeClient,
    *,
    department_id: str = "",
    user_id: str = "",
    page_size: int = 50,
) -> AsyncIterator[dict[str, Any]]:
    """Yield normalized WeChat friend records from the CRM.

    Walks ``POST /wechatweb/wxFriends/friend_list/es`` page by page until the
    upstream reports no more rows, mapping each record to the ``Customer``
    column shape (snake_case) plus a verbatim ``raw_json`` copy.

    Parameters
    ----------
    client:
        Authenticated :class:`YunkeClient`. Auth headers / 401 re-login are
        handled by the client; this scraper only concerns itself with the
        friend-list endpoint contract.
    department_id, user_id:
        Optional upstream filters (passed through to the request body as
        ``departmentId`` / ``userId``). Empty by default — scrape everything
        the logged-in account can see.
    page_size:
        Page size to request. Defaults to 50 (a sane middle ground between
        request count and per-response payload size for this endpoint).
    """
    async for record in paginate_by_page(
        client,
        FRIEND_LIST_PATH,
        page_size=page_size,
        parse_total=_parse_total,
        parse_rows=_parse_rows,
        build_body=_build_body,
        department_id=department_id,
        user_id=user_id,
    ):
        yield _normalize(record)


# ===========================================================================
# Message scraper (Task 1.6) — POST /wechatweb/msg/chats
# ===========================================================================


def _coerce_int_or_none(value: Any) -> int | None:
    """Coerce a possibly-string numeric value to ``int``; ``None`` preserved.

    Mirrors :func:`app.storage.repos._coerce_int` so the scraper can emit a
    ready-to-store ``type_raw`` int and check raw type codes (e.g. the ``-2``
    divider drop) without depending on the storage layer (the scraper has no
    DB dependency by design). A non-numeric value (``"text"`` / ``None`` /
    missing) returns ``None``.
    """
    if value is None:
        return None
    if isinstance(value, bool):
        return int(value)
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(value)
    if isinstance(value, str):
        s = value.strip()
        if not s:
            return None
        try:
            return int(s)
        except ValueError:
            return None
    return None


async def _fetch_msg_page(
    client: YunkeClient,
    *,
    wechat_id: str,
    friend_id: str,
    types: str,
    start: int,
    end_cursor: int,
) -> tuple[list[dict[str, Any]], int, bool]:
    """Fetch one backward page of ``msg/chats``.

    Returns the ``(rows, next_cursor, has_more)`` triple expected by
    :func:`paginate_by_cursor`. ``next_cursor`` is ``data.end - 1`` so the
    next call's window is ``(end-1, ...]`` — i.e. strictly older than the
    oldest message in this batch.
    """
    resp = await client.post(
        MSG_CHATS_PATH,
        data={
            "wechatId": wechat_id,
            "friendId": friend_id,
            "types": types,
            "start": str(start),
            "end": str(end_cursor),
        },
        headers={"content-type": _FORM_CONTENT_TYPE},
    )

    try:
        body = resp.json()
    except (AttributeError, ValueError):
        # Malformed envelope — treat as empty + no-more so the paginator stops.
        return [], end_cursor, False

    data = body.get("data") if isinstance(body, dict) else None
    if not isinstance(data, dict):
        return [], end_cursor, False

    rows = data.get("messages")
    if not isinstance(rows, list):
        rows = []

    end = data.get("end")
    next_cursor = (
        int(end) - 1 if isinstance(end, (int, float)) and not isinstance(end, bool)
        else end_cursor
    )
    has_more = bool(data.get("hasLast"))
    return rows, next_cursor, has_more


def _normalize_message(record: dict[str, Any]) -> dict[str, Any]:
    """Map one upstream ``msg/chats`` record to the ``Message`` column shape.

    Adds ``channel="wechat"`` and preserves the original record under
    ``raw_json``. Per design §4.3 the message type is carried as TWO fields:

    * ``type`` — the normalized enum label (``"text"`` / ``"unknown"`` / ...)
      via :func:`normalize_wechat_type`.
    * ``type_raw`` — the original upstream type code (``messageType`` falling
      back to ``type``), coerced to ``int | None`` so the raw code survives
      end-to-end without round-tripping through ``raw_json``.

    ``None``-safe: missing fields become ``None`` rather than raising, so a
    partial payload does not break the whole scrape. Value coercion
    (e.g. ``isSend``/``mine``/``isDel`` ``int`` -> ``bool``) is deferred to
    the persistence layer — here we carry the raw values through, except for
    ``is_send`` which is pulled from ``mine`` (falling back to ``isSend``)
    since that disambiguation is straightforward.
    """
    # `mine` and `isSend` carry the same semantic; the live capture uses both.
    # Prefer `mine`, fall back to `isSend`. Missing-key .get() returns None,
    # but `mine` may legitimately be the int 0 (falsy but present) — guard
    # explicitly with `is None`.
    mine = record.get("mine")
    is_send = mine if mine is not None else record.get("isSend")

    raw_type = record.get("messageType")
    if raw_type is None:
        raw_type = record.get("type")

    # `timestamp` is the canonical field; fall back to `createTime`.
    timestamp = record.get("timestamp")
    if timestamp is None:
        timestamp = record.get("createTime")

    return {
        "msg_id": record.get("msgSvrId"),
        "channel": "wechat",
        "owner_account_id": record.get("wechatId"),
        "talker_id": record.get("talker"),
        "ori_talker": record.get("oriTalker"),
        "room_id": record.get("roomid"),
        "is_send": is_send,
        "type": normalize_wechat_type(raw_type),
        "type_raw": _coerce_int_or_none(raw_type),
        "text": record.get("text"),
        "file_url": record.get("file"),
        "file_th_url": record.get("fileTh"),
        "file_hd_url": record.get("fileHd"),
        "md5": record.get("md5"),
        "md5_th": record.get("md5Th"),
        "md5_hd": record.get("md5Hd"),
        "duration": record.get("duration"),
        "transcription_text": record.get("transcriptionText"),
        "img_path": record.get("imgPath"),
        "original_image": record.get("originalImage"),
        "pic_type": record.get("picType"),
        "pay_msg_id": record.get("paymsgid"),
        "refer_msg": record.get("referMsg"),
        "sender_name": record.get("name"),
        "sender_head_url": record.get("headUrl"),
        "timestamp": timestamp,
        "is_del": record.get("isDel"),
        "version": record.get("version"),
        "raw_json": record,
    }


def _is_divider(record: dict[str, Any]) -> bool:
    """True iff this record is a date-divider (raw type code ``-2``).

    These are not real messages — the upstream inserts them as visual date
    bars in the chat UI — and must never be yielded or persisted. Per the
    Task 1.6 spec we drop a record when EITHER ``messageType`` or ``type`` is
    ``-2``; all other unknown codes pass through.

    The check is on the RAW code (not the normalized label) so the drop
    decision is robust against any future change to
    :func:`normalize_wechat_type`'s label vocabulary.
    """
    return _coerce_int_or_none(record.get("messageType")) == -2 or (
        _coerce_int_or_none(record.get("type")) == -2
    )


async def scrape_wechat_messages(
    client: YunkeClient,
    wechat_id: str,
    friend_id: str,
    start: int,
    end: int,
    *,
    types: str = "",
) -> AsyncIterator[dict[str, Any]]:
    """Yield normalized WeChat chat messages, paging backward by timestamp.

    Wraps ``POST /wechatweb/msg/chats`` (form-encoded body, NOT JSON). The
    endpoint pages BACKWARD: each response carries ``data.end`` (oldest
    timestamp in this batch) and ``data.hasLast`` (older messages exist?).
    The scraper advances the cursor to ``end - 1`` and repeats until
    ``hasLast`` is false, then stops.

    Date dividers (``type == -2``) are dropped — they are UI chrome, not real
    messages. All other type codes pass through (mapped to ``"unknown"`` when
    not in our small known-table) so no real message is silently lost.

    Parameters
    ----------
    client:
        Authenticated :class:`YunkeClient`. Auth headers / 401 re-login are
        handled by the client.
    wechat_id:
        Staff WeChat id (``wechatId`` form field) — maps to
        ``owner_account_id`` on each yielded record.
    friend_id:
        Friend ``wxid`` OR chatroom id (``friendId`` form field) — selects
        the conversation to pull.
    start, end:
        Time window in ms epoch. ``end`` is also the initial backward cursor;
        the first request queries the ``(start, end]`` window and the cursor
        walks toward ``start`` from there.
    types:
        Optional upstream message-type filter (passed through verbatim to the
        form body). Empty by default — pull all types.
    """
    async for record in paginate_by_cursor(
        lambda end_cursor: _fetch_msg_page(
            client,
            wechat_id=wechat_id,
            friend_id=friend_id,
            types=types,
            start=start,
            end_cursor=end_cursor,
        ),
        initial_cursor=end,
    ):
        if _is_divider(record):
            continue
        yield _normalize_message(record)