inbox/resolvers/load.py

import asyncio
import json
from typing import Any, Dict, List, Optional, Union

from services.auth import login_required
from services.core import get_my_followings, get_all_authors
from services.rediscache import redis
from services.schema import query
from validators.chat import Message, ChatPayload
from validators.member import ChatMember
from .chats import create_chat


async def get_unread_counter(chat_id: str, author_id: int) -> int:
    unread = await redis.execute("LLEN", f"chats/{chat_id}/unread/{author_id}")
    return unread or 0


# NOTE: not an API handler
async def load_messages(
    chat_id: str, limit: int = 5, offset: int = 0, ids: Optional[List[int]] = None
) -> List[Message]:
    """load :limit messages for :chat_id with :offset"""
    if ids is None:
        ids = []
    messages = []
    try:
        message_ids = [] + ids
        if limit:
            mids = (await redis.lrange(f"chats/{chat_id}/message_ids", offset, offset + limit)) or []
            message_ids += mids
        if message_ids:
            message_keys = [f"chats/{chat_id}/messages/{mid}" for mid in message_ids]
            messages = (await redis.mget(*message_keys)) or []
            messages = [json.loads(m) if isinstance(m, str) else m for m in messages]
            replies = []
            for m in messages:
                if m:
                    reply_to = m.get("replyTo")
                    if reply_to:
                        reply_to = int(reply_to)
                        if reply_to not in message_ids:
                            replies.append(reply_to)
            if replies:
                messages += await load_messages(chat_id, offset, limit, replies)
    except Exception:
        import traceback

        traceback.print_exc()
    return messages


@query.field("loadChats")
@login_required
async def load_chats(_, info, limit: int = 50, offset: int = 0) -> Dict[str, Union[List[Dict[str, Any]], None]]:
    """load :limit chats of current user with :offset"""
    author_id = info.context["author_id"]
    cids = (await redis.execute("SMEMBERS", f"chats_by_author/{author_id}")) or []
    members_online = (await redis.execute("SMEMBERS", "authors-online")) or []
    cids = list(cids)[offset : (offset + limit)]
    chats = []
    lock = asyncio.Lock()
    if len(cids) == 0:
        print(f"[resolvers.load] no chats for user with id={author_id}")
        r = await create_chat(None, info, members=[2])  # member with id = 2 is discours
        print(f"[resolvers.load] created chat: {r['chat']}")
        cids.append(r["chat"]["id"])
    all_authors: List[ChatMember] = await get_all_authors()
    authors = {a["id"]: a for a in all_authors}
    for cid in cids:
        async with lock:
            chat_str = await redis.execute("GET", f"chats/{cid}")
            print(f"[resolvers.load] redis GET by {cid}: {chat_str}")
            if chat_str:
                c: ChatPayload = json.loads(chat_str)
                c["messages"] = await load_messages(cid, 5, 0)
                c["unread"] = await get_unread_counter(cid, author_id)
                member_ids = c["members"].copy()
                c["members"] = []
                for member_id in member_ids:
                    a = authors.get(member_id)
                    if a:
                        a["online"] = a.get("id") in members_online
                        c["members"].append(a)
                chats.append(c)
    return {"chats": chats, "error": None}


@query.field("loadMessagesBy")
@login_required
async def load_messages_by(_, info, by, limit: int = 10, offset: int = 0):
    """load :limit messages of :chat_id with :offset"""
    author_id = info.context["author_id"]
    user_chats = (await redis.execute("SMEMBERS", "chats_by_author/" + str(author_id))) or []
    user_chats = [c for c in user_chats]
    if user_chats:
        messages = []
        by_chat = by.get("chat")
        if by_chat in user_chats:
            chat = await redis.execute("GET", f"chats/{by_chat}")
            if not chat:
                return {"messages": [], "error": "chat not exist"}
            # everyone's messages in filtered chat
            messages = await load_messages(by_chat, limit, offset)
        return {
            "messages": sorted(
                [m for m in messages if m.get("created_at")],
                key=lambda m: m.get("created_at"),
            ),
            "error": None,
        }
    else:
        return {"error": "Cannot access messages of this chat"}


@query.field("loadRecipients")
async def load_recipients(_, _info, limit=50, offset=0):
    """load possible chat participants"""
    onliners = (await redis.execute("SMEMBERS", "authors-online")) or []
    r = []
    all_authors: List[ChatMember] = await get_all_authors()
    my_followings = await get_my_followings()
    if len(my_followings) < limit:
        my_followings = my_followings + all_authors[0 : limit - len(my_followings)]
    for a in my_followings:
        a["online"] = a["id"] in onliners
        r.append(a)

    # NOTE: maybe sort members here

    print(f"[resolvers.load] loadRecipients found {len(r)} members")

    return {"members": r, "error": None}