yunke/yunke-server/app/storage/repos.py

"""Idempotent upsert repositories bridging scraper dicts to ORM tables.

Every public helper takes an ``AsyncSession`` plus one-or-more scraper-shaped
dicts (the snake_case records that ``app.scraper.*`` yields) and writes them
to the corresponding table via **dialect-aware Core upsert**, so the same call
path is correct on:

* **in-memory SQLite** — used by the test suite (``sqlite+aiosqlite``); uses
  ``sqlalchemy.dialects.sqlite.insert(...).on_conflict_do_update(...)``, and
* **MySQL** — production (``mysql+asyncmy``); uses
  ``sqlalchemy.dialects.mysql.insert(...).on_duplicate_key_update(...)``.

The dialect is detected at write time via ``session.bind.name``, so callers
do not need to know which backend they target.

Type coercion
-------------
Upstream CRM values are not always typed the way the columns want them. The
following coercions are applied when building column values from the input
dict (see :func:`_coerce_bool`, :func:`_coerce_int`):

* Booleans (``is_block`` / ``is_send`` / ``is_del`` / ``is_on_job`` /
  ``is_upload``): ``1`` / ``"1"`` / ``True`` -> ``True``; ``0`` / ``"0"`` /
  ``False`` -> ``False``; ``None`` preserved.
* ``gender``: ``"0"`` / ``"1"`` / ``"2"`` (str) -> ``int`` when not ``None``.
* ``type_raw`` (Message): the raw upstream type code, coerced to
  ``int | None``. ``Message.type`` is a normalized label string and is
  stored as-is (no int coercion).
* PhoneCall int fields (``call_duration`` / ``start_call_time`` /
  ``call_status`` / ``incoming_call`` / ``qc_status``): coerced to
  ``int | None``. The voice-call scraper already parses these to ints
  (incl. ``startTime`` datetime -> ms epoch); the repo coercion is a
  defensive net for hand-built records that arrive as numeric strings.
* Timestamps (``last_chat_time`` / ``timestamp`` / ``last_update`` / ...):
  passed through as-is (ms-epoch ``int`` or ``None``).
* ``raw_json``: the original record dict, preserved verbatim. When the input
  already carries a ``raw_json`` key (scraper output does), that value wins;
  otherwise the input dict itself is stored.

Sharding note
-------------
``app.db.models.message_table_name`` exists for the future per-month sharded
``message_<channel>_<YYYYMM>`` layout. For M1, the :class:`Message` model maps
to a single physical ``message`` table and :func:`upsert_messages` writes to
that one table. Multi-table routing is a LATER task (YAGNI for M1).
"""

from __future__ import annotations

from typing import Any

import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase

from app.db.models import (
    ChatGroup,
    Customer,
    Department,
    GroupMember,
    Media,
    Message,
    PhoneCall,
    Sms,
    Staff,
    WechatAccount,
    WhatsappAccount,
)

__all__ = [
    "upsert_chat_group",
    "upsert_customer",
    "upsert_customers",
    "upsert_department",
    "upsert_media",
    "upsert_messages",
    "upsert_phone_call",
    "upsert_sms",
    "upsert_staff",
    "upsert_staff_from_account",
    "upsert_voice_call",
    "upsert_wechat_account",
    "upsert_whatsapp_account",
]


# ---------------------------------------------------------------------------
# Type coercion helpers (module-level so they are unit-testable)
# ---------------------------------------------------------------------------

# Columns whose source values are int-as-bool and must land as actual bool.
_BOOL_FIELDS: frozenset[str] = frozenset(
    {"is_block", "is_send", "is_del", "is_on_job", "is_upload"}
)

# Columns whose source values may arrive as numeric strings.
_INT_FIELDS: frozenset[str] = frozenset(
    {
        "gender",
        "type_raw",
        "pic_type",
        "friend_state",
        "delete_flag",
        "direction",
        "incoming_call",
        "call_status",
        "call_duration",
        "start_call_time",
        "qc_status",
        "create_time",
    }
)


def _coerce_bool(value: Any) -> bool | None:
    """Coerce int/str/bool CRM bool fields to a real bool.

    ``None`` is preserved. ``1``/``"1"``/``True`` -> ``True``;
    ``0``/``"0"``/``False`` -> ``False``. Common textual forms
    (``"true"``/``"false"``/``"yes"``/``"no"``) are also accepted; any other
    non-empty string falls back to truthiness so a stray value never raises.
    """
    if value is None:
        return None
    if isinstance(value, bool):
        return value
    if isinstance(value, int):
        return value != 0
    if isinstance(value, str):
        s = value.strip().lower()
        if s in ("1", "true", "yes", "y", "t"):
            return True
        if s in ("0", "false", "no", "n", "f", ""):
            return False
        # Last-resort: try numeric compare, else truthiness of the raw string.
        try:
            return int(s) != 0
        except ValueError:
            return bool(s)
    # Floats / other: Python truthiness.
    return bool(value)


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

    ``"2"`` -> ``2``; ``2`` -> ``2``; ``0``/``"0"`` -> ``0``; ``None`` ->
    ``None``. A non-numeric string (e.g. ``"text"``) returns ``None`` — the
    caller (e.g. Message.type_raw) treats that as "no value".
    """
    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


# ---------------------------------------------------------------------------
# Dialect-aware Core upsert
# ---------------------------------------------------------------------------


def _pk_col_names(model: type[DeclarativeBase]) -> list[str]:
    """Return the names of the model's primary-key columns."""
    return [c.name for c in model.__table__.primary_key.columns]


def _build_column_values(
    model: type[DeclarativeBase], rec: dict[str, Any]
) -> dict[str, Any]:
    """Pick rec keys that match model columns, applying type coercion.

    * Only keys that correspond to real columns on ``model`` are kept (extra
      upstream keys are dropped, since they would cause Core insert() to raise
      on unknown columns).
    * Boolean / int-as-str fields are coerced via :func:`_coerce_bool` /
      :func:`_coerce_int`.
    * ``raw_json`` defaults to the entire input dict when not explicitly
      provided, so the verbatim upstream record is always recoverable.
    """
    col_names = {c.name for c in model.__table__.columns}
    values: dict[str, Any] = {}
    for col_name in col_names:
        if col_name in rec:
            v = rec[col_name]
            if col_name in _BOOL_FIELDS:
                v = _coerce_bool(v)
            elif col_name in _INT_FIELDS:
                v = _coerce_int(v)
            # Some upstream fields are lists (e.g. Customer.deptRoute is a
            # list of dept ids) but map to scalar String columns — join so the
            # bind succeeds and the data is still recoverable (raw_json keeps
            # the original list).
            if isinstance(v, list):
                v = ";".join(str(x) for x in v)
            values[col_name] = v
    # Preserve the verbatim source record. If the scraper already set raw_json
    # (it does), that wins; otherwise the input dict itself is the source of
    # truth.
    if "raw_json" in col_names and "raw_json" not in values:
        values["raw_json"] = rec
    return values


def _build_upsert_stmt(
    model: type[DeclarativeBase], values: dict[str, Any], dialect_name: str
) -> sa.Insert:
    """Build a dialect-specific idempotent Core upsert for one row.

    On conflict the statement updates **every non-PK column** from the
    incoming row, so re-running a scrape yields one final row per PK with the
    freshest values (counts, version, names, raw_json, updated_at, ...).

    Parameters
    ----------
    model:
        The ORM model class to upsert into.
    values:
        Column-name -> value dict (output of :func:`_build_column_values`).
    dialect_name:
        One of ``"sqlite"`` / ``"mysql"`` (the two backends this project
        targets). Any other dialect raises ``ValueError``.
    """
    pk_cols = _pk_col_names(model)
    # Non-PK columns to refresh on conflict (volatile fields).
    update_col_names = [c.name for c in model.__table__.columns if c.name not in pk_cols]

    # NB: access ``excluded`` / ``inserted`` on the BARE insert (before
    # ``.values(...)``) and reuse that accessor when building the SET dict.
    # If the accessor is taken on the post-``.values()`` statement, the
    # ``excluded.<col>`` references do not bind to the supplied values at
    # compile time and the ON CONFLICT DO UPDATE silently fails to refresh
    # the row (a SQLAlchemy quirk; verified on 2.0.x).
    if dialect_name == "sqlite":
        from sqlalchemy.dialects.sqlite import insert as sqlite_insert

        bare = sqlite_insert(model)
        excluded = bare.excluded
        set_ = {name: excluded[name] for name in update_col_names}
        return bare.values(**values).on_conflict_do_update(
            index_elements=pk_cols, set_=set_
        )

    if dialect_name == "mysql":
        from sqlalchemy.dialects.mysql import insert as mysql_insert

        bare = mysql_insert(model)
        inserted = bare.inserted
        update_dict = {name: inserted[name] for name in update_col_names}
        return bare.values(**values).on_duplicate_key_update(**update_dict)

    raise ValueError(
        f"Unsupported dialect for idempotent upsert: {dialect_name!r}. "
        "Only 'sqlite' (tests) and 'mysql' (prod) are supported."
    )


async def _dialect_name(session: AsyncSession) -> str:
    """Return the dialect name of the session's bind (e.g. 'sqlite'/'mysql').

    Uses ``session.bind.name`` (the AsyncEngine); falls back to
    ``session.get_bind().dialect.name`` if ``bind`` is not set on the session.
    """
    bind = session.bind
    if bind is not None and hasattr(bind, "name"):
        return bind.name  # type: ignore[no-any-return]
    return session.get_bind().dialect.name


async def _upsert_one(session: AsyncSession, model: type[DeclarativeBase], rec: dict[str, Any]) -> None:
    """Upsert a single record into ``model``."""
    values = _build_column_values(model, rec)
    # If the record carries no values for any column (e.g. empty dict), there
    # is nothing to write — skip the round-trip. PK-only rows still go through
    # because their PK is in `values`.
    if not values:
        return
    dialect = await _dialect_name(session)
    stmt = _build_upsert_stmt(model, values, dialect)
    await session.execute(stmt)


# ---------------------------------------------------------------------------
# Public upsert helpers
# ---------------------------------------------------------------------------


async def upsert_customer(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert one :class:`Customer` (PK = ``friend_id``).

    Same ``friend_id`` twice -> one row, with non-PK fields refreshed from the
    later record.
    """
    await _upsert_one(session, Customer, rec)


async def upsert_customers(session: AsyncSession, recs: list[dict[str, Any]]) -> int:
    """Batch upsert for :class:`Customer`. Returns the number of records processed.

    Idempotent: duplicate ``friend_id`` values within the batch collapse to a
    single row (the last write wins) at the DB layer.
    """
    for rec in recs:
        await _upsert_one(session, Customer, rec)
    return len(recs)


async def upsert_messages(session: AsyncSession, recs: list[dict[str, Any]]) -> int:
    """Batch idempotent upsert for :class:`Message` (PK = ``msg_id``).

    Applies the documented coercions (``is_send``/``is_del`` -> bool,
    ``type_raw`` -> int). ``Message.type`` is the normalized label string and
    is stored as-is. Returns the number of records processed (input length);
    the DB collapses duplicate ``msg_id`` values to one row.

    .. note::
       For M1 the ``Message`` model maps to a single physical ``message``
       table. Real per-month sharding across
       ``message_<channel>_<YYYYMM>`` tables (see
       :func:`app.db.models.message_table_name`) is a LATER task; do not
       implement multi-table routing here.
    """
    for rec in recs:
        await _upsert_one(session, Message, rec)
    return len(recs)


# ---------------------------------------------------------------------------
# Department upsert (Task 2.0)
# ---------------------------------------------------------------------------


async def upsert_department(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert one :class:`Department` (PK = ``dept_id``).

    Same ``dept_id`` twice -> one row, with non-PK fields (``parent_id`` /
    ``name`` / ``path`` / ``raw_json``) refreshed from the later record.
    """
    await _upsert_one(session, Department, rec)


# ---------------------------------------------------------------------------
# WechatAccount + embedded Staff (best-effort)
# ---------------------------------------------------------------------------


async def upsert_staff(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert one :class:`Staff` row from a staff-shaped dict.

    PK = ``user_id``. Same ``user_id`` twice -> one row, with non-PK fields
    refreshed. Applies the documented coercions (``is_on_job`` -> bool).

    This is the primitive staff upsert; the account-scoped helper
    :func:`upsert_staff_from_account` (which extracts the staff fields from a
    wechat-account record) delegates here, and :func:`upsert_wechat_account`
    calls that helper so an account upsert automatically refreshes its
    embedded staff row too.
    """
    await _upsert_one(session, Staff, rec)


async def upsert_wechat_account(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert a :class:`WechatAccount` AND its embedded Staff row.

    PK = ``wx_id``. The companion :func:`upsert_staff_from_account` covers the
    embedded staff row (``user_id`` + ``user_name``); this helper calls it
    automatically so an account scrape that lands accounts also lands the
    owning staff — callers do not need to upsert staff separately.

    The staff upsert is best-effort: an account record with no ``user_id``
    (e.g. an orphaned / unassigned wechat number) still writes the account
    row; the staff upsert silently no-ops.
    """
    await _upsert_one(session, WechatAccount, rec)
    await upsert_staff_from_account(session, rec)


async def upsert_staff_from_account(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Best-effort upsert of the Staff row embedded in a wechat-account record.

    Maps the account's ``user_id`` + ``user_name`` (plus ``is_on_job`` /
    ``phone`` / ``sale_phone`` when present) onto :class:`Staff`
    (PK = ``user_id``) and delegates to :func:`upsert_staff`. Silently no-ops
    when the record has no ``user_id`` — there is nothing to key the row on.

    Per Task 2.1, dept linkage for staff is inferred ELSEWHERE (from
    customers' ``deptRoute``); staff with no customers have no dept. This
    helper therefore does not populate ``dept_id`` — that is a later concern.
    """
    user_id = rec.get("user_id")
    if not user_id:
        return
    staff_rec: dict[str, Any] = {
        "user_id": user_id,
        "name": rec.get("user_name"),
        "raw_json": rec,
    }
    # Forward any Staff-relevant fields the account record happens to carry.
    # ``phone`` and ``sale_phone`` are usually absent on the wxUser payload
    # (they come from the WhatsApp account payload); carry them through when
    # present so a future shared caller does not silently drop them.
    for key in ("is_on_job", "phone", "sale_phone"):
        if key in rec:
            staff_rec[key] = rec[key]
    await upsert_staff(session, staff_rec)


async def upsert_whatsapp_account(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert a :class:`WhatsappAccount` AND its embedded Staff row.

    PK = ``uid`` (of the form ``<phone>@s.whatsapp.net``). Mirrors
    :func:`upsert_wechat_account`: the companion :func:`upsert_staff_from_account`
    covers the embedded staff row (``user_id`` + ``user_name``); this helper
    calls it automatically so an account scrape that lands accounts also lands
    the owning staff — callers do not need to upsert staff separately.

    The WA account payload (unlike the WeChat one) DOES carry ``sale_phone``,
    so the embedded staff row picks it up via the shared staff-from-account
    helper. The staff upsert is best-effort: an account record with no
    ``user_id`` still writes the account row; the staff upsert silently
    no-ops.
    """
    await _upsert_one(session, WhatsappAccount, rec)
    await upsert_staff_from_account(session, rec)


# ---------------------------------------------------------------------------
# ChatGroup + GroupMember upsert (M2)
# ---------------------------------------------------------------------------


async def upsert_chat_group(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert one :class:`ChatGroup` AND its :class:`GroupMember` rows.

    PK = ``room_id``. The group row is upserted first; then, for every id in
    the record's ``member_ids_list`` (the scraper-parsed, empties-dropped
    member wxid list), a :class:`GroupMember` row is upserted against the
    composite PK ``(channel, room_id, member_id)``. Re-scraping a group
    therefore refreshes the group row and re-asserts its members without
    duplication (idempotent).

    The ``member_ids_list`` key is the parsed list shape produced by
    :func:`app.scraper.wechat._normalize_chat_group`; it is NOT a
    ``ChatGroup`` column, so it is dropped by :func:`_build_column_values`
    when writing the group row (the raw semicolon-joined string is preserved
    on ``ChatGroup.member_ids``).

    Per-member ``nickname`` / ``avatar`` / ``display_name`` are left ``None``
    here — Phase-0 does not yet resolve them (they need a separate per-member
    API; see data dictionary §9.6). The repo writes a minimal member row keyed
    by the composite PK so membership is captured now and enriched later.
    """
    await _upsert_one(session, ChatGroup, rec)

    channel = rec.get("channel") or "wechat"
    room_id = rec.get("room_id")
    # No room_id -> nothing to key member rows on (defensive; the scraper
    # always sets room_id). Skip the member loop rather than raising so a
    # malformed record still lands its (PK-less) group upsert attempt above.
    if not room_id:
        return

    member_ids: list[Any] = list(rec.get("member_ids_list") or [])
    for member_id in member_ids:
        # Defensive: never insert an empty/falsy member id (would collide on
        # an empty-PK row). The scraper already drops empties, but a hand-
        # built record could slip one through.
        if not member_id:
            continue
        member_rec: dict[str, Any] = {
            "channel": channel,
            "room_id": room_id,
            "member_id": member_id,
        }
        await _upsert_one(session, GroupMember, member_rec)


# ---------------------------------------------------------------------------
# PhoneCall (voice_call) upsert (M2, Task 2.4)
# ---------------------------------------------------------------------------


async def upsert_voice_call(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert one :class:`PhoneCall` (PK = ``call_id``).

    Same ``call_id`` twice -> one row, with non-PK fields refreshed from the
    later record. Applies the documented int coercions (``call_duration`` /
    ``start_call_time`` / ``call_status`` / ``incoming_call`` -> int when they
    arrive as numeric strings; the voice-call scraper already emits ints, but
    a hand-built record may not). ``is_upload`` is a bool field and goes
    through :func:`_coerce_bool` like the other bool columns.

    The voice-call channel (``channel="voice_call"``) and the phone-call
    channel (``channel="phone_call"``) share the ``PhoneCall`` table; the
    per-row ``channel`` value disambiguates them.
    """
    await _upsert_one(session, PhoneCall, rec)


# ---------------------------------------------------------------------------
# PhoneCall (phone_call) + Sms upsert (M4)
# ---------------------------------------------------------------------------


async def upsert_phone_call(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert one :class:`PhoneCall` from a phone-call record (M4).

    Same ``call_id`` twice -> one row, with non-PK fields refreshed from the
    later record. PK = ``call_id`` (the source ``id``). The phone-call channel
    (``channel="phone_call"``) shares the :class:`PhoneCall` table with the
    voice-call channel; the per-row ``channel`` value disambiguates them.

    Applies the documented int coercions (``call_duration`` /
    ``start_call_time`` / ``call_status`` / ``incoming_call`` / ``qc_status``
    -> int when they arrive as numeric strings; the phone-call scraper already
    emits ints, but a hand-built record may not). ``is_upload`` is a bool
    field and goes through :func:`_coerce_bool` like the other bool columns.
    """
    await _upsert_one(session, PhoneCall, rec)


async def upsert_sms(session: AsyncSession, rec: dict[str, Any]) -> None:
    """Idempotently upsert one :class:`Sms` (PK = ``sms_id``) (M4).

    Same ``sms_id`` twice -> one row, with non-PK fields refreshed from the
    later record. Applies the documented int coercions (``direction`` /
    ``create_time`` -> int when they arrive as numeric strings; the SMS
    scraper already emits ints, but a hand-built record may not). The
    ``channel`` column defaults to ``"sms"`` at the model level; the scraper
    carries it explicitly so the upsert is independent of the column-default
    application order.
    """
    await _upsert_one(session, Sms, rec)


# ---------------------------------------------------------------------------
# Media upsert (M2, Task 2.6)
# ---------------------------------------------------------------------------

# Download artifacts that must NOT be clobbered with NULL when a later upsert
# (e.g. a re-scrape that did not re-download) omits them. COALESCE'd against the
# existing row so a genuine new value still wins.
_MEDIA_PRESERVE_IF_NULL: frozenset[str] = frozenset(
    {"local_path", "size_bytes", "mime_type"}
)


async def upsert_media(
    session: AsyncSession,
    *,
    media_hash: str,
    channel: str,
    media_type: str | None = None,
    original_url: str | None = None,
    original_th_url: str | None = None,
    original_hd_url: str | None = None,
    local_path: str | None = None,
    size_bytes: int | None = None,
    mime_type: str | None = None,
    duration: int | None = None,
    raw_json: dict[str, Any] | None = None,
) -> None:
    """Idempotently upsert one :class:`Media` row (PK = ``media_hash``).

    Same ``media_hash`` twice -> one row. Non-PK columns are refreshed from the
    incoming record on conflict, EXCEPT the volatile download artifacts
    (``local_path`` / ``size_bytes`` / ``mime_type``): those are written with
    ``COALESCE(incoming, existing)`` so that a re-scrape which did NOT re-download
    (and therefore passes ``None``) leaves the previously downloaded values
    intact. A genuine new value (non-``None``) still overrides the old one.

    The ``media_hash`` is the cross-message dedup key: the source ``md5`` when
    the message carried one, else the sha256 of the downloaded bytes (see
    :mod:`app.scraper.media`).
    """
    values: dict[str, Any] = {
        "media_hash": media_hash,
        "channel": channel,
        "media_type": media_type,
        "original_url": original_url,
        "original_th_url": original_th_url,
        "original_hd_url": original_hd_url,
        "local_path": local_path,
        "size_bytes": size_bytes,
        "mime_type": mime_type,
        "duration": duration,
        "raw_json": raw_json,
    }

    dialect = await _dialect_name(session)
    pk_cols = _pk_col_names(Media)
    update_cols = [
        c.name for c in Media.__table__.columns if c.name not in pk_cols
    ]
    media_cols = Media.__table__.c

    def _build_set(other_cols: Any) -> dict[str, Any]:
        """Map each updated column to its conflict-set expression.

        Volatile download artifacts use ``COALESCE(incoming, existing)``; every
        other column is refreshed directly from the incoming row.
        """
        set_: dict[str, Any] = {}
        for name in update_cols:
            if name in _MEDIA_PRESERVE_IF_NULL:
                set_[name] = sa.func.coalesce(other_cols[name], media_cols[name])
            else:
                set_[name] = other_cols[name]
        return set_

    if dialect == "sqlite":
        from sqlalchemy.dialects.sqlite import insert as sqlite_insert

        bare = sqlite_insert(Media)
        set_ = _build_set(bare.excluded)
        stmt = bare.values(**values).on_conflict_do_update(
            index_elements=pk_cols, set_=set_
        )
    elif dialect == "mysql":
        from sqlalchemy.dialects.mysql import insert as mysql_insert

        bare = mysql_insert(Media)
        set_ = _build_set(bare.inserted)
        stmt = bare.values(**values).on_duplicate_key_update(**set_)
    else:
        raise ValueError(
            f"Unsupported dialect for upsert_media: {dialect!r}. "
            "Only 'sqlite' (tests) and 'mysql' (prod) are supported."
        )

    await session.execute(stmt)