from __future__ import annotations import json import logging import os import tempfile import threading import time from dataclasses import dataclass, field from collections import OrderedDict logger = logging.getLogger(__name__) DEFAULT_PERSIST_DIR = os.path.join(tempfile.gettempdir(), "yuxi_qqbot_known_users") @dataclass class UserRecord: user_id: str username: str = "" first_seen: float = field(default_factory=time.time) last_seen: float = field(default_factory=time.time) message_count: int = 0 chat_types: set[str] = field(default_factory=set) metadata: dict = field(default_factory=dict) def touch(self, username: str = "", chat_type: str = "") -> None: self.last_seen = time.time() self.message_count += 1 if username: self.username = username if chat_type: self.chat_types.add(chat_type) def to_dict(self) -> dict: return { "user_id": self.user_id, "username": self.username, "first_seen": self.first_seen, "last_seen": self.last_seen, "message_count": self.message_count, "chat_types": list(self.chat_types), "metadata": self.metadata, } @classmethod def from_dict(cls, data: dict) -> UserRecord: return cls( user_id=data["user_id"], username=data.get("username", ""), first_seen=data.get("first_seen", time.time()), last_seen=data.get("last_seen", time.time()), message_count=data.get("message_count", 0), chat_types=set(data.get("chat_types", [])), metadata=data.get("metadata", {}), ) class KnownUserTracker: def __init__( self, app_id: str, max_users: int = 10000, persist_dir: str | None = None, persist_interval_s: int = 300, ): self._app_id = app_id self._max_users = max_users self._persist_dir = persist_dir or DEFAULT_PERSIST_DIR self._persist_path = os.path.join(self._persist_dir, f"{app_id}_users.json") self._persist_interval = persist_interval_s self._lock = threading.Lock() self._users: OrderedDict[str, UserRecord] = OrderedDict() self._last_persist: float = 0 self._dirty = False self._load_from_disk() def record(self, user_id: str, username: str = "", chat_type: str = "") -> UserRecord: with self._lock: if user_id in self._users: self._users.move_to_end(user_id) record = self._users[user_id] record.touch(username, chat_type) else: record = UserRecord(user_id=user_id, username=username) record.touch(username, chat_type) self._users[user_id] = record self._users.move_to_end(user_id) while len(self._users) > self._max_users: self._users.popitem(last=False) self._dirty = True self._maybe_persist() return record def is_known(self, user_id: str) -> bool: with self._lock: return user_id in self._users def get(self, user_id: str) -> UserRecord | None: with self._lock: return self._users.get(user_id) def remove(self, user_id: str) -> bool: with self._lock: if user_id in self._users: del self._users[user_id] self._dirty = True return True return False @property def count(self) -> int: with self._lock: return len(self._users) def persist(self) -> bool: with self._lock: if not self._dirty: return True return self._do_persist() def _maybe_persist(self) -> None: now = time.time() if now - self._last_persist < self._persist_interval: return if not self._dirty: return self._do_persist() def _do_persist(self) -> bool: try: os.makedirs(self._persist_dir, exist_ok=True) data = { "app_id": self._app_id, "updated_at": time.time(), "users": {uid: u.to_dict() for uid, u in self._users.items()}, } tmp_path = self._persist_path + ".tmp" with open(tmp_path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False) os.replace(tmp_path, self._persist_path) self._last_persist = time.time() self._dirty = False logger.debug("KnownUserTracker: persisted %d users", len(self._users)) return True except OSError: logger.exception("KnownUserTracker: persist failed") return False def _load_from_disk(self) -> None: try: if not os.path.exists(self._persist_path): return with open(self._persist_path, encoding="utf-8") as f: data = json.load(f) users_data = data.get("users", {}) for uid, udata in users_data.items(): self._users[uid] = UserRecord.from_dict(udata) logger.info("KnownUserTracker: loaded %d users from disk", len(self._users)) except (OSError, json.JSONDecodeError): logger.exception("KnownUserTracker: failed to load from disk") def clear(self) -> None: with self._lock: self._users.clear() self._dirty = True self.persist() def get_recent_users(self, limit: int = 50) -> list[UserRecord]: with self._lock: return list(reversed(self._users.values()))[:limit]