ForcePilot/backend/package/yuxi/channels/adapters/yuanbao/directory.py

from __future__ import annotations

import aiohttp
from dataclasses import dataclass, field

from yuxi.utils.logging_config import logger


@dataclass
class PeerInfo:
    open_id: str
    username: str = ""
    avatar_url: str = ""
    display_name: str = ""


@dataclass
class GroupInfo:
    group_open_id: str
    name: str = ""
    avatar_url: str = ""
    member_count: int = 0
    bot_is_member: bool = False


@dataclass
class DirectoryResult:
    peers: list[PeerInfo] = field(default_factory=list)
    groups: list[GroupInfo] = field(default_factory=list)


async def get_user_info(
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
    user_id: str,
) -> PeerInfo | None:
    headers = {"Authorization": f"Bearer {token}"}
    url = f"{api_base}/api/v1/user/{user_id}"

    try:
        async with http_client.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=10)) as resp:
            if resp.status == 200:
                data = await resp.json()
                return PeerInfo(
                    open_id=data.get("open_id", user_id),
                    username=data.get("username", ""),
                    avatar_url=data.get("avatar_url", ""),
                    display_name=data.get("display_name", ""),
                )
            else:
                logger.debug(f"[Yuanbao] get_user_info returned {resp.status} for {user_id}")
                return None
    except Exception as e:
        logger.warning(f"[Yuanbao] Failed to get user info for {user_id}: {e}")
        return None


async def get_group_list(
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
) -> list[GroupInfo]:
    headers = {"Authorization": f"Bearer {token}"}
    url = f"{api_base}/api/v1/bot/groups"

    try:
        async with http_client.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=10)) as resp:
            if resp.status == 200:
                data = await resp.json()
                groups = []
                for g in data.get("groups", []):
                    groups.append(
                        GroupInfo(
                            group_open_id=g.get("group_open_id", ""),
                            name=g.get("group_name", ""),
                            avatar_url=g.get("avatar_url", ""),
                            member_count=g.get("member_count", 0),
                            bot_is_member=g.get("bot_is_member", False),
                        )
                    )
                logger.info(f"[Yuanbao] Fetched {len(groups)} groups")
                return groups
            else:
                logger.warning(f"[Yuanbao] get_group_list returned {resp.status}")
                return []
    except Exception as e:
        logger.warning(f"[Yuanbao] Failed to get group list: {e}")
        return []


async def get_group_detail(
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
    group_open_id: str,
) -> GroupInfo | None:
    headers = {"Authorization": f"Bearer {token}"}
    url = f"{api_base}/api/v1/bot/groups/{group_open_id}"

    try:
        async with http_client.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=10)) as resp:
            if resp.status == 200:
                data = await resp.json()
                return GroupInfo(
                    group_open_id=data.get("group_open_id", group_open_id),
                    name=data.get("group_name", ""),
                    avatar_url=data.get("avatar_url", ""),
                    member_count=data.get("member_count", 0),
                    bot_is_member=data.get("bot_is_member", False),
                )
            else:
                logger.debug(f"[Yuanbao] get_group_detail returned {resp.status} for {group_open_id}")
                return None
    except Exception as e:
        logger.warning(f"[Yuanbao] Failed to get group detail for {group_open_id}: {e}")
        return None


async def get_group_members(
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
    group_open_id: str,
) -> list[PeerInfo]:
    headers = {"Authorization": f"Bearer {token}"}
    url = f"{api_base}/api/v1/bot/groups/{group_open_id}/members"

    try:
        async with http_client.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=10)) as resp:
            if resp.status == 200:
                data = await resp.json()
                members = []
                for m in data.get("members", []):
                    members.append(
                        PeerInfo(
                            open_id=m.get("open_id", ""),
                            username=m.get("username", ""),
                            avatar_url=m.get("avatar_url", ""),
                            display_name=m.get("display_name", ""),
                        )
                    )
                logger.info(f"[Yuanbao] Fetched {len(members)} members from group {group_open_id}")
                return members
            else:
                logger.warning(f"[Yuanbao] get_group_members returned {resp.status} for {group_open_id}")
                return []
    except Exception as e:
        logger.warning(f"[Yuanbao] Failed to get group members for {group_open_id}: {e}")
        return []


async def list_directory(
    api_base: str,
    token: str,
    http_client: aiohttp.ClientSession,
    user_ids: list[str] | None = None,
) -> DirectoryResult:
    result = DirectoryResult()

    result.groups = await get_group_list(api_base, token, http_client)

    if user_ids:
        for user_id in user_ids:
            peer = await get_user_info(api_base, token, http_client, user_id)
            if peer:
                result.peers.append(peer)

    return result