ForcePilot/backend/package/yuxi/channels/adapters/qqbot/known_users.py

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]