zhenxun_bot/zhenxun/services/hot_query_cache.py

from __future__ import annotations

import asyncio
from collections.abc import Iterable
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Literal

from tortoise.functions import Count

from zhenxun.services.cache.bounded_ttl import BoundedTTLCache


@dataclass(frozen=True, slots=True)
class GroupMemberSnapshot:
    id: int
    user_id: str
    user_name: str
    group_id: str
    user_join_time: datetime | None
    uid: int | None
    platform: str | None


def _member_cache_sizeof(members: tuple[GroupMemberSnapshot, ...]) -> int:
    size = 0
    for member in members:
        size += 96
        size += len(member.user_id) + len(member.user_name) + len(member.group_id)
        size += len(member.platform or "")
    return size


_GROUP_MEMBER_CACHE = BoundedTTLCache[str, tuple[GroupMemberSnapshot, ...]](
    "hot_group_info_users",
    ttl_seconds=45,
    max_items=512,
    max_total_bytes=32 * 1024 * 1024,
    sizeof=_member_cache_sizeof,
)
_GROUP_USER_IDS_CACHE = BoundedTTLCache[str, tuple[str, ...]](
    "hot_group_info_user_ids",
    ttl_seconds=45,
    max_items=2048,
)
_GROUP_MEMBER_BY_ID_CACHE = BoundedTTLCache[str, tuple[GroupMemberSnapshot | None]](
    "hot_group_info_user_by_id",
    ttl_seconds=45,
    max_items=50000,
)
_USER_GROUP_CACHE = BoundedTTLCache[str, tuple[str, ...]](
    "hot_group_info_user_groups",
    ttl_seconds=45,
    max_items=4096,
)
_USER_NAME_CACHE = BoundedTTLCache[str, str](
    "hot_group_info_user_names",
    ttl_seconds=45,
    max_items=20000,
)
_CHAT_RANK_CACHE = BoundedTTLCache[str, tuple[tuple[str, int], ...]](
    "hot_chat_history_rank",
    ttl_seconds=20,
    max_items=512,
)
_CHAT_FIRST_MSG_CACHE = BoundedTTLCache[str, tuple[datetime | None]](
    "hot_chat_history_first_msg",
    ttl_seconds=300,
    max_items=2048,
)
_STATISTICS_COUNT_CACHE = BoundedTTLCache[str, tuple[tuple[str, int], ...]](
    "hot_statistics_plugin_counts",
    ttl_seconds=20,
    max_items=512,
)

_GROUP_MEMBER_LOCKS: dict[str, asyncio.Lock] = {}
_GROUP_USER_IDS_LOCKS: dict[str, asyncio.Lock] = {}
_USER_GROUP_LOCKS: dict[str, asyncio.Lock] = {}
_CHAT_RANK_LOCKS: dict[str, asyncio.Lock] = {}
_CHAT_FIRST_MSG_LOCKS: dict[str, asyncio.Lock] = {}
_STATISTICS_LOCKS: dict[str, asyncio.Lock] = {}
_MAX_LOCK_POOL_SIZE = 4096


def _get_lock(pool: dict[str, asyncio.Lock], key: str) -> asyncio.Lock:
    lock = pool.get(key)
    if lock is None:
        if len(pool) >= _MAX_LOCK_POOL_SIZE:
            for old_key, old_lock in list(pool.items()):
                if not old_lock.locked():
                    pool.pop(old_key, None)
                    break
        lock = asyncio.Lock()
        pool[key] = lock
    return lock


def _normalize_id(value: object) -> str:
    return str(value or "")


def _normalize_ids(values: Iterable[object] | None) -> tuple[str, ...] | None:
    if values is None:
        return None
    return tuple(dict.fromkeys(v for value in values if (v := _normalize_id(value))))


async def get_group_members(
    group_id: str | int | None,
) -> tuple[GroupMemberSnapshot, ...]:
    """Return lightweight group-member snapshots with a short runtime TTL."""
    group_key = _normalize_id(group_id)
    if not group_key:
        return ()

    cached = await _GROUP_MEMBER_CACHE.get(group_key)
    if cached is not None:
        return cached

    lock = _get_lock(_GROUP_MEMBER_LOCKS, group_key)
    async with lock:
        cached = await _GROUP_MEMBER_CACHE.get(group_key)
        if cached is not None:
            return cached

        from zhenxun.models.group_member_info import GroupInfoUser

        rows = await GroupInfoUser.filter(group_id=group_key).values_list(
            "id",
            "user_id",
            "user_name",
            "user_join_time",
            "uid",
            "platform",
        )
        members = tuple(
            GroupMemberSnapshot(
                id=int(row[0] or 0),
                user_id=str(row[1] or ""),
                user_name=str(row[2] or ""),
                group_id=group_key,
                user_join_time=row[3],
                uid=int(row[4]) if row[4] is not None else None,
                platform=str(row[5]) if row[5] else None,
            )
            for row in rows
            if row[1]
        )
        await _GROUP_MEMBER_CACHE.set(group_key, members)
        await _GROUP_USER_IDS_CACHE.set(
            group_key, tuple(member.user_id for member in members)
        )
        return members


async def get_group_member_map(
    group_id: str | int | None,
    user_ids: Iterable[object] | None = None,
) -> dict[str, GroupMemberSnapshot]:
    group_key = _normalize_id(group_id)
    if not group_key:
        return {}
    wanted = _normalize_ids(user_ids)
    if wanted is not None and not wanted:
        return {}
    if wanted is None:
        members = await get_group_members(group_key)
        return {member.user_id: member for member in members}

    cached_members = await _GROUP_MEMBER_CACHE.get(group_key)
    if cached_members is not None:
        wanted_set = set(wanted)
        return {
            member.user_id: member
            for member in cached_members
            if member.user_id in wanted_set
        }

    result: dict[str, GroupMemberSnapshot] = {}
    missing: list[str] = []
    for user_id in wanted:
        cache_key = f"{group_key}:{user_id}"
        cached = await _GROUP_MEMBER_BY_ID_CACHE.get(cache_key)
        if cached is None:
            missing.append(user_id)
        else:
            member = cached[0]
            if member is not None:
                result[user_id] = member

    if missing:
        from zhenxun.models.group_member_info import GroupInfoUser

        rows = await GroupInfoUser.filter(
            group_id=group_key, user_id__in=missing
        ).values_list(
            "id",
            "user_id",
            "user_name",
            "user_join_time",
            "uid",
            "platform",
        )
        found: set[str] = set()
        for row in rows:
            if not row[1]:
                continue
            member = GroupMemberSnapshot(
                id=int(row[0] or 0),
                user_id=str(row[1] or ""),
                user_name=str(row[2] or ""),
                group_id=group_key,
                user_join_time=row[3],
                uid=int(row[4]) if row[4] is not None else None,
                platform=str(row[5]) if row[5] else None,
            )
            result[member.user_id] = member
            found.add(member.user_id)
            await _GROUP_MEMBER_BY_ID_CACHE.set(
                f"{group_key}:{member.user_id}", (member,)
            )
        for user_id in missing:
            if user_id not in found:
                await _GROUP_MEMBER_BY_ID_CACHE.set(f"{group_key}:{user_id}", (None,))
    return result


async def get_group_member(
    group_id: str | int | None,
    user_id: str | int | None,
) -> GroupMemberSnapshot | None:
    user_key = _normalize_id(user_id)
    if not user_key:
        return None
    return (await get_group_member_map(group_id, [user_key])).get(user_key)


async def get_group_user_ids(group_id: str | int | None) -> set[str]:
    group_key = _normalize_id(group_id)
    if not group_key:
        return set()
    cached = await _GROUP_USER_IDS_CACHE.get(group_key)
    if cached is not None:
        return set(cached)

    cached_members = await _GROUP_MEMBER_CACHE.get(group_key)
    if cached_members is not None:
        user_ids = tuple(member.user_id for member in cached_members)
        await _GROUP_USER_IDS_CACHE.set(group_key, user_ids)
        return set(user_ids)

    lock = _get_lock(_GROUP_USER_IDS_LOCKS, group_key)
    async with lock:
        cached = await _GROUP_USER_IDS_CACHE.get(group_key)
        if cached is not None:
            return set(cached)

        from zhenxun.models.group_member_info import GroupInfoUser

        rows = await GroupInfoUser.filter(group_id=group_key).values_list(
            "user_id", flat=True
        )
        user_ids = tuple(str(user_id) for user_id in rows if user_id)
        await _GROUP_USER_IDS_CACHE.set(group_key, user_ids)
        return set(user_ids)


async def get_user_group_ids(user_id: str | int | None) -> list[str]:
    user_key = _normalize_id(user_id)
    if not user_key:
        return []

    cached = await _USER_GROUP_CACHE.get(user_key)
    if cached is not None:
        return list(cached)

    lock = _get_lock(_USER_GROUP_LOCKS, user_key)
    async with lock:
        cached = await _USER_GROUP_CACHE.get(user_key)
        if cached is not None:
            return list(cached)

        from zhenxun.models.group_member_info import GroupInfoUser

        rows = await GroupInfoUser.filter(user_id=user_key).values_list(
            "group_id", flat=True
        )
        group_ids = tuple(str(group_id) for group_id in rows if group_id)
        await _USER_GROUP_CACHE.set(user_key, group_ids)
        return list(group_ids)


async def get_member_names(
    user_ids: Iterable[object],
    group_id: str | int | None = None,
) -> dict[str, str]:
    user_keys = _normalize_ids(user_ids) or ()
    if not user_keys:
        return {}
    if group_id:
        members = await get_group_member_map(group_id, user_keys)
        return {user_id: members[user_id].user_name for user_id in members}

    result: dict[str, str] = {}
    missing: list[str] = []
    for user_id in user_keys:
        cached = await _USER_NAME_CACHE.get(user_id)
        if cached is None:
            missing.append(user_id)
        else:
            result[user_id] = cached

    if missing:
        from zhenxun.models.group_member_info import GroupInfoUser

        rows = await GroupInfoUser.filter(user_id__in=missing).values_list(
            "user_id", "user_name"
        )
        for user_id, user_name in rows:
            user_key = str(user_id)
            if user_key not in result:
                result[user_key] = str(user_name or "")
        for user_id in missing:
            await _USER_NAME_CACHE.set(user_id, result.get(user_id, ""))
    return result


async def get_member_name(
    user_id: str | int | None,
    group_id: str | int | None = None,
) -> str | None:
    user_key = _normalize_id(user_id)
    if not user_key:
        return None
    return (await get_member_names([user_key], group_id)).get(user_key) or None


async def invalidate_group_members(
    group_id: str | int | None = None,
    user_ids: Iterable[object] | None = None,
) -> None:
    if group_id is None:
        await _GROUP_MEMBER_CACHE.clear()
        await _GROUP_USER_IDS_CACHE.clear()
        await _GROUP_MEMBER_BY_ID_CACHE.clear()
        return
    group_key = _normalize_id(group_id)
    await _GROUP_MEMBER_CACHE.delete(group_key)
    await _GROUP_USER_IDS_CACHE.delete(group_key)
    normalized_ids = _normalize_ids(user_ids)
    if normalized_ids is None:
        await _GROUP_MEMBER_BY_ID_CACHE.clear()
        return
    for user_id in normalized_ids:
        await _GROUP_MEMBER_BY_ID_CACHE.delete(f"{group_key}:{user_id}")


async def invalidate_member_names(user_ids: Iterable[object] | None = None) -> None:
    if user_ids is None:
        await _USER_NAME_CACHE.clear()
        await _USER_GROUP_CACHE.clear()
        return
    for user_id in _normalize_ids(user_ids) or ():
        await _USER_NAME_CACHE.delete(user_id)
        await _USER_GROUP_CACHE.delete(user_id)


def _datetime_key(value: datetime | None) -> str:
    return value.isoformat(" ", timespec="seconds") if value else ""


def _date_scope_key(date_scope: tuple[datetime, datetime] | None) -> str:
    if not date_scope:
        return ""
    end_bucket = int(date_scope[1].timestamp() // 20)
    return f"{_datetime_key(date_scope[0])}..bucket:{end_bucket}"


async def get_chat_history_rank_cached(
    model: Any,
    gid: str | None,
    limit: int = 10,
    order: str = "DESC",
    date_scope: tuple[datetime, datetime] | None = None,
) -> list[tuple[str, int]]:
    key = f"{gid or '*'}:{limit}:{order}:{_date_scope_key(date_scope)}"
    cached = await _CHAT_RANK_CACHE.get(key)
    if cached is not None:
        return list(cached)

    lock = _get_lock(_CHAT_RANK_LOCKS, key)
    async with lock:
        cached = await _CHAT_RANK_CACHE.get(key)
        if cached is not None:
            return list(cached)

        order_prefix = "-" if order == "DESC" else ""
        query: Any = model.filter(group_id=gid) if gid else model
        if date_scope:
            filter_scope = (
                date_scope[0].isoformat(" "),
                date_scope[1].isoformat(" "),
            )
            query = query.filter(create_time__range=filter_scope)
        rows = (
            await query.annotate(count=Count("user_id"))
            .order_by(f"{order_prefix}count")
            .group_by("user_id")
            .limit(limit)
            .values_list("user_id", "count")
        )
        result = tuple((str(user_id), int(count)) for user_id, count in rows)
        await _CHAT_RANK_CACHE.set(key, result)
        return list(result)


async def get_chat_history_first_msg_datetime_cached(
    model: Any,
    group_id: str | None,
) -> datetime | None:
    key = group_id or "*"
    cached = await _CHAT_FIRST_MSG_CACHE.get(key)
    if cached is not None:
        return cached[0]

    lock = _get_lock(_CHAT_FIRST_MSG_LOCKS, key)
    async with lock:
        cached = await _CHAT_FIRST_MSG_CACHE.get(key)
        if cached is not None:
            return cached[0]

        query: Any = model.filter(group_id=group_id) if group_id else model.all()
        message = await query.order_by("create_time").first()
        result = getattr(message, "create_time", None) if message else None
        await _CHAT_FIRST_MSG_CACHE.set(key, (result,))
        return result


async def get_statistics_plugin_counts_cached(
    scope: Literal["global", "user", "group"],
    *,
    plugin_name: str | None,
    start_time: datetime | None,
    user_id: str | None = None,
    group_id: str | None = None,
) -> list[tuple[str, int]]:
    key = (
        f"{scope}:{plugin_name or ''}:{_datetime_key(start_time)}:"
        f"{user_id or ''}:{group_id or ''}"
    )
    cached = await _STATISTICS_COUNT_CACHE.get(key)
    if cached is not None:
        return list(cached)

    lock = _get_lock(_STATISTICS_LOCKS, key)
    async with lock:
        cached = await _STATISTICS_COUNT_CACHE.get(key)
        if cached is not None:
            return list(cached)

        from zhenxun.models.statistics import Statistics

        query: Any = Statistics
        if scope == "user":
            query = Statistics.filter(user_id=user_id)
            if group_id:
                query = query.filter(group_id=group_id)
        elif scope == "group":
            query = Statistics.filter(group_id=group_id)
        if plugin_name:
            query = query.filter(plugin_name=plugin_name)
        if start_time:
            query = query.filter(create_time__gte=start_time)
        rows = (
            await query.annotate(count=Count("id"))
            .group_by("plugin_name")
            .values_list("plugin_name", "count")
        )
        result = tuple((str(plugin), int(count)) for plugin, count in rows)
        await _STATISTICS_COUNT_CACHE.set(key, result)
        return list(result)