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

"""WeChat + organization scrapers.

Endpoints covered 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).
* :func:`scrape_departments` — wraps ``POST /usercenter/dept/tree`` and
  yields normalized :class:`app.db.models.Department` dicts. Single request,
  no pagination; the tree (flat or nested) is flattened and the materialized
  ``path`` is reconstructed from parent links (Task 2.0).
* :func:`scrape_wechat_accounts` — wraps ``POST /wechatweb/wxUser/list`` and
  yields normalized :class:`app.db.models.WechatAccount` dicts (plus the
  embedded ``userId``/``userName`` staff pair). Page-number pagination
  (Task 2.1).

All yield dicts (not ORM). Persistence is the caller's job (the
:mod:`app.storage.repos` module bridges these dicts to the ORM tables). Each
dict carries:

* snake_case keys matching the target ORM columns (see :func:`_normalize`
  for friends, :func:`_normalize_message` for messages,
  :func:`_normalize_department` for departments, and
  :func:`_normalize_wxuser_account` for accounts),
* ``channel == "wechat"`` (only on communication-channel entities; the
  organization entities ``Department`` / ``WechatAccount`` carry no channel),
* ``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_departments",
    "scrape_wechat_accounts",
    "scrape_wechat_friends",
    "scrape_wechat_messages",
]

FRIEND_LIST_PATH = "/wechatweb/wxFriends/friend_list/es"
MSG_CHATS_PATH = "/wechatweb/msg/chats"
DEPT_TREE_PATH = "/usercenter/dept/tree"
WXUSER_LIST_PATH = "/wechatweb/wxUser/list"

# 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)


# ===========================================================================
# Department tree scraper (Task 2.0) — POST /usercenter/dept/tree
# ===========================================================================


def _extract_dept_node_list(resp: Any) -> list[dict[str, Any]]:
    """Extract the raw node list from a ``dept/tree`` response, defensively.

    The envelope's ``data`` may be:

    * a list of nodes (the common case), OR
    * an object wrapping the node list under ``children`` (or ``list`` /
      ``tree`` / ``data`` — be tolerant of wrapper-key variance).

    Returns ``[]`` on any malformed shape so the scraper yields nothing rather
    than raising: a broken upstream should not abort the whole sync. Nodes are
    returned AS-IS (possibly still nested via ``children``); flattening is the
    next step (:func:`_flatten_dept_tree`).
    """
    try:
        body = resp.json()
    except (AttributeError, ValueError):
        return []
    data = body.get("data") if isinstance(body, dict) else None
    if isinstance(data, list):
        return data
    if isinstance(data, dict):
        # Walk the plausible wrapper keys in priority order. ``children`` is
        # the documented yunke tree convention; the others are defensive.
        for key in ("children", "list", "tree", "data"):
            val = data.get(key)
            if isinstance(val, list):
                return val
    return []


def _flatten_dept_tree(nodes: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Flatten a (possibly nested) dept-tree node list depth-first.

    Each node MAY carry a ``children`` array (the nested shape); when absent,
    the tree is already flat and parents are carried via ``pId``. Both shapes
    flatten to the same depth-first list.

    Each emitted node is a shallow copy of the source node with the
    ``children`` key stripped, so the ``raw_json`` we store per-dept does not
    carry a recursive subtree (only that node's own attributes). Non-dict
    entries are dropped defensively.
    """
    flat: list[dict[str, Any]] = []
    # Iterative DFS via a stack. Nodes are pushed in reverse so the first
    # child is popped first (preserves source order in the output).
    stack: list[dict[str, Any]] = [
        n for n in reversed(nodes) if isinstance(n, dict)
    ]
    while stack:
        node = stack.pop()
        children = node.get("children")
        # Shallow copy without ``children`` so raw_json stored per-dept is clean.
        flat_node = {k: v for k, v in node.items() if k != "children"}
        flat.append(flat_node)
        if isinstance(children, list) and children:
            stack.extend(
                reversed([c for c in children if isinstance(c, dict)])
            )
    return flat


def _build_dept_paths(nodes: list[dict[str, Any]]) -> dict[str, str]:
    """Compute the materialized path (ancestor names joined by ``/``) per node.

    Returns ``{id: path}``. The path is built by walking up the ``pId`` chain;
    the walk terminates when it reaches a node whose ``pId`` is falsy
    (``None`` / ``""`` / ``0`` / ``"0"``) OR whose parent is not in the tree
    (handles the common yunke convention where root carries ``pId="0"`` and no
    node with ``id="0"`` exists). A visited set caps the walk so a cyclic
    ``pId`` chain cannot loop forever.

    A node whose own ``id`` is missing is skipped (no path entry) — the caller
    treats a missing path as empty string.
    """
    by_id: dict[str, dict[str, Any]] = {}
    for n in nodes:
        nid = n.get("id")
        # First occurrence wins so duplicate ids cannot shadow the canonical
        # node (defensive against an upstream that emits a node twice).
        if nid is not None and nid not in by_id:
            by_id[nid] = n

    paths: dict[str, str] = {}
    for nid, node in by_id.items():
        parts: list[str] = []
        seen: set[str] = set()
        cur: dict[str, Any] | None = node
        while cur is not None:
            cid = cur.get("id")
            if cid is None or cid in seen:
                break
            seen.add(cid)
            name = cur.get("name")
            if name is not None:
                parts.append(name)
            pid = cur.get("pId")
            # Falsy pId OR unknown parent terminates the walk.
            cur = by_id.get(pid) if pid else None
        paths[nid] = "/".join(reversed(parts))
    return paths


def _normalize_department(
    node: dict[str, Any], path: str
) -> dict[str, Any]:
    """Map one dept-tree node to the ``Department`` column shape.

    ``parent_id`` is carried verbatim from ``pId`` — the scraper does NOT
    synthesize ``NULL`` for root nodes; downstream code treats falsy
    ``parent_id`` (``"0"`` / ``""`` / ``None``) as "no parent" consistently
    with the materialized-path builder. ``path`` is the precomputed
    slash-joined ancestor-name chain from :func:`_build_dept_paths`.
    """
    return {
        "dept_id": node.get("id"),
        "parent_id": node.get("pId"),
        "name": node.get("name"),
        "path": path,
        "raw_json": node,
    }


async def scrape_departments(
    client: YunkeClient,
) -> AsyncIterator[dict[str, Any]]:
    """Yield normalized department records from ``POST /usercenter/dept/tree``.

    Single request (no pagination). The endpoint follows yunke's tree
    convention: nodes carry ``id`` / ``pId`` / ``name`` and MAY be delivered
    either as a flat list (parents via ``pId``) or as a nested tree (each node
    optionally carrying ``children``). Both shapes are flattened to one flat
    list of records, each carrying the materialized ``path`` (ancestor names
    joined with ``/``).

    Defensive envelope: ``data`` may be the node list directly OR an object
    wrapping the node list under ``children``; both are unwrapped.

    The async-generator signature mirrors the other scrapers so callers can
    consume every entity stream with the same ``async for`` pattern, even
    though this endpoint always resolves in one round-trip.
    """
    resp = await client.post(DEPT_TREE_PATH, json={})
    nodes = _extract_dept_node_list(resp)
    flat = _flatten_dept_tree(nodes)
    paths = _build_dept_paths(flat)
    for node in flat:
        nid = node.get("id")
        path = paths.get(nid, "") if nid is not None else ""
        yield _normalize_department(node, path)


# ===========================================================================
# WeChat staff-account scraper (Task 2.1) — POST /wechatweb/wxUser/list
# ===========================================================================


def _parse_wxuser_total(resp: Any) -> int | None:
    """Read the declared row count from the ``wxUser/list`` envelope.

    Shape: ``data.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
    total = data.get("total")
    return int(total) if isinstance(total, (int, float)) and not isinstance(
        total, bool
    ) else None


def _parse_wxuser_rows(resp: Any) -> list[dict[str, Any]]:
    """Extract the account list from the ``wxUser/list`` envelope.

    Shape: ``data.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 []
    rows = data.get("data")
    return rows if isinstance(rows, list) else []


def _build_wxuser_body(page_num: int, page_size: int) -> dict[str, Any]:
    """Build the ``wxUser/list`` request body for a given page.

    Unlike ``friend_list/es``, the ``wxUser/list`` body carries no
    staff/dept filter parameters in the live capture — just the page slice.
    """
    return {"pageNum": page_num, "pageSize": page_size}


def _normalize_wxuser_account(record: dict[str, Any]) -> dict[str, Any]:
    """Map one upstream ``wxUser/list`` record to the ``WechatAccount`` shape.

    Carries the embedded staff pair (``user_id`` <- ``userId``,
    ``user_name`` <- ``userName``) so the persistence layer can upsert the
    owning :class:`app.db.models.Staff` row alongside the account. Adds
    ``raw_json`` (the verbatim upstream record). ``None``-safe: missing fields
    become ``None`` rather than raising, so partial payloads do not break the
    scrape.
    """
    return {
        "wx_id": record.get("wechatId"),
        "nick_name": record.get("nickName"),
        "alias": record.get("alias"),
        "phone": record.get("phone"),
        "user_id": record.get("userId"),
        "user_name": record.get("userName"),
        "contacts_num": record.get("contactsNum"),
        "group_num": record.get("groupNum"),
        "send_num": record.get("sendNum"),
        "receive_num": record.get("receiveNum"),
        "is_on_job": record.get("isOnJob"),
        "db_size": record.get("dbSize"),
        "online_status": record.get("onlineStatus"),
        "last_update": record.get("lastUpdate"),
        "version": record.get("version"),
        "raw_json": record,
    }


async def scrape_wechat_accounts(
    client: YunkeClient,
    *,
    page_size: int = 50,
) -> AsyncIterator[dict[str, Any]]:
    """Yield normalized WeChat staff-account records from ``wxUser/list``.

    Walks ``POST /wechatweb/wxUser/list`` page by page until the upstream
    reports no more rows (empty ``data.data`` array — the primary stop signal),
    mapping each record to the ``WechatAccount`` column shape plus a verbatim
    ``raw_json`` copy.

    The envelope differs from ``friend_list/es``: there is **no**
    ``basicResult`` wrapper. The pagination fields sit directly under ``data``
    (``data.data`` / ``data.total`` / ``data.pageNum`` / ``data.pageSize``).

    Parameters
    ----------
    client:
        Authenticated :class:`YunkeClient`. Auth headers / 401 re-login are
        handled by the client.
    page_size:
        Page size to request. Defaults to 50 (matches the live capture's
        observed page size for this endpoint).
    """
    async for record in paginate_by_page(
        client,
        WXUSER_LIST_PATH,
        page_size=page_size,
        parse_total=_parse_wxuser_total,
        parse_rows=_parse_wxuser_rows,
        build_body=_build_wxuser_body,
    ):
        yield _normalize_wxuser_account(record)