Sync phone contacts and enable chat deletion

pymax-worker/src/server.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import base64
+import json
 import mimetypes
 import os
 import pathlib
@@ -18,12 +19,14 @@ from typing import Any
 import aiohttp
 from aiohttp import web
 from pymax import Client, ExtraConfig, File, Photo, SyncOverrides, Video
+from pymax.types import ContactInfo
 
 
 PORT = int(os.environ.get("PORT", "3002"))
 PHONE_NUMBER = os.environ.get("PYMAX_PHONE_NUMBER") or os.environ.get("QMAX_MAX_PHONE_NUMBER") or ""
 SESSION_DIR = pathlib.Path(os.environ.get("PYMAX_SESSION_DIR", "/data/pymax-session"))
 SESSION_NAME = os.environ.get("PYMAX_SESSION_NAME", "session.db")
+PHONE_CONTACT_IDS_FILE = SESSION_DIR / "phone-contact-ids.json"
 CHAT_FETCH_LIMIT = int(os.environ.get("PYMAX_CHAT_FETCH_LIMIT", "350"))
 HISTORY_LIMIT = int(os.environ.get("PYMAX_HISTORY_LIMIT", "80"))
 SEND_ROOTS = [pathlib.Path(p) for p in os.environ.get("PYMAX_SEND_ROOTS", "/qmax-data:/tmp").split(":") if p]
@@ -56,6 +59,11 @@ def normalize_phone(phone: str) -> str:
     return phone.strip()
 
 
+def phone_key(phone: Any) -> str:
+    normalized = normalize_phone(str(phone or ""))
+    return "".join(ch for ch in normalized if ch.isdigit())
+
+
 def is_expired_session_error(error: str | None) -> bool:
     text = (error or "").lower()
     return any(
@@ -253,7 +261,13 @@ def user_avatar_url(user: Any) -> str | None:
     return first_text(data.get("base_url"), data.get("base_raw_url"))
 
 
-def contact_result(contact: Any, me_id: int, external_chat_id: Any | None = None) -> dict[str, Any] | None:
+def contact_result(
+    contact: Any,
+    me_id: int,
+    external_chat_id: Any | None = None,
+    *,
+    is_saved_contact: bool = False,
+) -> dict[str, Any] | None:
     data = dump_model(contact)
     contact_id = coerce_int(data.get("id") or getattr(contact, "id", None))
     if contact_id is None or contact_id <= 0 or contact_id == me_id:
@@ -265,6 +279,7 @@ def contact_result(contact: Any, me_id: int, external_chat_id: Any | None = None
         "avatarUrl": user_avatar_url(contact),
         "phoneNumber": first_text(data.get("phone")),
         "status": first_text(data.get("status"), data.get("description")),
+        "isSavedContact": is_saved_contact,
     }
 
 
@@ -523,10 +538,6 @@ def normalize_chat_kind(chat: Any) -> str:
     return "MaxDialog"
 
 
-def is_direct_dialog(chat: Any) -> bool:
-    return value_name(getattr(chat, "type", "")).upper() == "DIALOG"
-
-
 def normalize_attachment_kind(att: Any) -> str:
     data = dump_model(att)
     raw_type = value_name(data.get("type") or data.get("_type") or getattr(att, "type", "")).lower()
@@ -935,27 +946,151 @@ async def contacts(_request: web.Request) -> web.Response:
     me_id = get_me_user_id(client)
     if me_id is None:
         return json_response({"success": False, "error": "MAX profile has no user id."}, status=503)
-
-    by_user_id: dict[str, dict[str, Any]] = {}
-    for contact in client.contacts:
-        normalized = contact_result(contact, me_id) if contact is not None else None
-        if normalized is not None:
-            by_user_id[normalized["userId"]] = normalized
-
-    chats = await fetch_chat_pages(client, CHAT_FETCH_LIMIT)
-    direct_chats = [chat for chat in chats if is_direct_dialog(chat)]
-    user_map = await build_user_map(client, direct_chats)
-    for chat in direct_chats:
-        contact = chat_other_user(chat, user_map, me_id)
-        normalized = contact_result(contact, me_id, getattr(chat, "id", None)) if contact is not None else None
-        if normalized is not None:
-            by_user_id.setdefault(normalized["userId"], normalized)
-
-    result = list(by_user_id.values())
+    matched_ids = load_phone_contact_ids()
+    if not matched_ids:
+        return json_response([])
+    users = await client.get_users(sorted(matched_ids))
+    saved_ids = {
+        coerce_int(dump_model(contact).get("id") or getattr(contact, "id", None))
+        for contact in client.contacts
+        if contact is not None
+    }
+    result = [
+        normalized
+        for contact in users
+        if (normalized := contact_result(
+            contact,
+            me_id,
+            is_saved_contact=coerce_int(dump_model(contact).get("id") or getattr(contact, "id", None)) in saved_ids,
+        )) is not None
+    ]
     result.sort(key=lambda item: item["displayName"].casefold())
     return json_response(result)
 
 
+def load_phone_contact_ids() -> set[int]:
+    try:
+        values = json.loads(PHONE_CONTACT_IDS_FILE.read_text(encoding="utf-8"))
+    except (OSError, ValueError, TypeError):
+        return set()
+    return {value for item in values if (value := coerce_int(item)) is not None and value > 0}
+
+
+def save_phone_contact_ids(contact_ids: set[int]) -> None:
+    SESSION_DIR.mkdir(parents=True, exist_ok=True)
+    part_path = PHONE_CONTACT_IDS_FILE.with_suffix(".json.part")
+    part_path.write_text(json.dumps(sorted(contact_ids)), encoding="utf-8")
+    part_path.replace(PHONE_CONTACT_IDS_FILE)
+
+
+def append_saved_contacts(client: Client, contacts: list[Any]) -> None:
+    saved = client.contacts
+    saved_ids = {
+        coerce_int(dump_model(contact).get("id") or getattr(contact, "id", None))
+        for contact in saved
+        if contact is not None
+    }
+    for contact in contacts:
+        contact_id = coerce_int(dump_model(contact).get("id") or getattr(contact, "id", None))
+        if contact_id is not None and contact_id not in saved_ids:
+            saved.append(contact)
+            saved_ids.add(contact_id)
+
+
+@route_errors
+async def contacts_import(request: web.Request) -> web.Response:
+    data = await read_json(request)
+    raw_contacts = data.get("contacts") if isinstance(data.get("contacts"), list) else []
+    if not raw_contacts:
+        return json_response({"success": False, "error": "contacts are required"}, status=400)
+    if len(raw_contacts) > 5000:
+        return json_response({"success": False, "error": "contacts limit is 5000"}, status=400)
+
+    contacts_to_import: list[ContactInfo] = []
+    requested_phone_keys: set[str] = set()
+    for item in raw_contacts:
+        if not isinstance(item, dict):
+            continue
+        phone = normalize_phone(str(item.get("phoneNumber") or item.get("phone") or "").strip())
+        first_name = str(item.get("firstName") or item.get("name") or phone).strip()
+        last_name = str(item.get("lastName") or "").strip() or None
+        key = phone_key(phone)
+        if phone and key and key not in requested_phone_keys:
+            requested_phone_keys.add(key)
+            contacts_to_import.append(ContactInfo(phone=phone, first_name=first_name or phone, last_name=last_name))
+    if not contacts_to_import:
+        return json_response({"success": False, "error": "contacts contain no phone numbers"}, status=400)
+
+    client = await runtime.get_client()
+    imported = list(await client.import_contacts(contacts_to_import) or [])
+    append_saved_contacts(client, imported)
+
+    # PyMax's SYNC response may omit contacts that were already saved in MAX.
+    # Merge those cached contacts back by phone so a repeated phone-book sync
+    # produces the same intersection instead of replacing it with an empty set.
+    matched_by_id: dict[int, Any] = {}
+    for contact in imported:
+        contact_id = coerce_int(dump_model(contact).get("id") or getattr(contact, "id", None))
+        if contact_id is not None and contact_id > 0:
+            matched_by_id[contact_id] = contact
+    for contact in client.contacts:
+        contact_data = dump_model(contact)
+        contact_id = coerce_int(contact_data.get("id") or getattr(contact, "id", None))
+        if (
+            contact_id is not None
+            and contact_id > 0
+            and phone_key(contact_data.get("phone") or getattr(contact, "phone", None)) in requested_phone_keys
+        ):
+            matched_by_id[contact_id] = contact
+
+    matched = list(matched_by_id.values())
+    imported_ids = {
+        contact_id
+        for contact in matched
+        if (contact_id := coerce_int(dump_model(contact).get("id") or getattr(contact, "id", None))) is not None
+        and contact_id > 0
+    }
+    save_phone_contact_ids(imported_ids)
+    me_id = get_me_user_id(client)
+    result = [
+        normalized
+        for contact in matched
+        if me_id is not None
+        if (normalized := contact_result(contact, me_id, is_saved_contact=True)) is not None
+    ]
+    return json_response(result)
+
+
+@route_errors
+async def contacts_add(request: web.Request) -> web.Response:
+    data = await read_json(request)
+    user_id = coerce_int(data.get("userId"))
+    if user_id is None or user_id <= 0:
+        return json_response({"success": False, "error": "valid userId is required"}, status=400)
+    client = await runtime.get_client()
+    contact = await client.add_contact(user_id)
+    append_saved_contacts(client, [contact])
+    me_id = get_me_user_id(client)
+    normalized = contact_result(contact, me_id, is_saved_contact=True) if me_id is not None else None
+    return json_response(normalized or {"success": False, "error": "contact could not be normalized"})
+
+
+@route_errors
+async def contacts_remove(request: web.Request) -> web.Response:
+    data = await read_json(request)
+    user_id = coerce_int(data.get("userId"))
+    if user_id is None or user_id <= 0:
+        return json_response({"success": False, "error": "valid userId is required"}, status=400)
+    client = await runtime.get_client()
+    await client.remove_contact(user_id)
+    client.contacts[:] = [
+        contact
+        for contact in client.contacts
+        if contact is not None and coerce_int(dump_model(contact).get("id") or getattr(contact, "id", None)) != user_id
+    ]
+    return json_response({"success": True, "error": None})
+
+
 @route_errors
 async def chat_history(request: web.Request) -> web.Response:
     data = await read_json(request)
@@ -1035,6 +1170,20 @@ async def disabled_action(_request: web.Request) -> web.Response:
     return json_response({"success": False, "error": "Not implemented in PyMax POC worker."})
 
 
+@route_errors
+async def chat_delete(request: web.Request) -> web.Response:
+    data = await read_json(request)
+    client = await runtime.get_client()
+    chat_id = parse_chat_id(data.get("externalChatId") or data.get("chatUrl"))
+    chat = await client.get_chat(chat_id)
+    await client.delete_chat(
+        chat_id,
+        last_event_time=getattr(chat, "last_event_time", None),
+        for_all=False,
+    )
+    return json_response({"success": True, "error": None})
+
+
 @route_errors
 async def send_text(request: web.Request) -> web.Response:
     data = await read_json(request)
@@ -1083,13 +1232,16 @@ def create_app() -> web.Application:
     app.router.add_get("/snapshot", snapshot)
     app.router.add_get("/updates", updates)
     app.router.add_get("/contacts", contacts)
+    app.router.add_post("/contacts/import", contacts_import)
+    app.router.add_post("/contacts/add", contacts_add)
+    app.router.add_post("/contacts/remove", contacts_remove)
     app.router.add_post("/chat/history", chat_history)
     app.router.add_post("/chat/resolve-url", resolve_chat_url)
     app.router.add_post("/channels/search", channels_search)
     app.router.add_post("/channels/join", channels_join)
     app.router.add_post("/chat/presence", chat_presence)
     app.router.add_post("/chat/clear-history", disabled_action)
-    app.router.add_post("/chat/delete", disabled_action)
+    app.router.add_post("/chat/delete", chat_delete)
     app.router.add_post("/send/text", send_text)
     app.router.add_post("/send/attachment", send_attachment)
     app.router.add_post("/message/edit", disabled_action)