from __future__ import annotations import asyncio import hashlib import re from collections import defaultdict from datetime import UTC, datetime, timedelta from secrets import token_hex from typing import Any from zoneinfo import ZoneInfo, ZoneInfoNotFoundError from pymongo import ASCENDING, DESCENDING from pymongo.errors import DuplicateKeyError from wbb import db SOURCE_CHECKIN = "checkin" SOURCE_ACTIVITY = "activity" SOURCE_UPVOTE = "upvote" SOURCE_ADMIN = "admin_adjustment" SOURCE_GIVEAWAY_ENTRY = "giveaway_entry" SOURCE_GIVEAWAY_REFUND = "giveaway_refund" SOURCE_GIVEAWAY_PARTICIPATION = "giveaway_participation" SOURCE_GIVEAWAY_WINNER = "giveaway_winner" POINT_SOURCES = { SOURCE_CHECKIN, SOURCE_ACTIVITY, SOURCE_UPVOTE, SOURCE_ADMIN, SOURCE_GIVEAWAY_ENTRY, SOURCE_GIVEAWAY_REFUND, SOURCE_GIVEAWAY_PARTICIPATION, SOURCE_GIVEAWAY_WINNER, } DEFAULT_POINT_RULES: dict[str, Any] = { "enabled": False, "timezone": "Asia/Shanghai", "checkin_enabled": True, "checkin_button_enabled": False, "checkin_points": 10, "activity_enabled": True, "activity_points": 1, "activity_cooldown_seconds": 300, "activity_daily_cap": 10, "upvote_enabled": True, "upvote_points": 2, "upvote_pair_cooldown_seconds": 86400, "upvote_daily_cap": 10, } accountsdb = db.point_accounts transactionsdb = db.point_transactions rulesdb = db.point_rules rate_limitsdb = db.point_rate_limits _index_lock = asyncio.Lock() _indexes_ready = False _account_locks: defaultdict[tuple[int, int], asyncio.Lock] = defaultdict(asyncio.Lock) _rate_locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock) class PointsError(RuntimeError): pass class InsufficientPoints(PointsError): pass def utc_now() -> datetime: return datetime.now(UTC) def _as_utc(value: datetime) -> datetime: if value.tzinfo is None: return value.replace(tzinfo=UTC) return value.astimezone(UTC) def _local_day(now: datetime, timezone_name: str) -> str: try: zone = ZoneInfo(timezone_name) except ZoneInfoNotFoundError: zone = ZoneInfo("Asia/Shanghai") return _as_utc(now).astimezone(zone).date().isoformat() def _safe_int(value: Any, default: int, minimum: int, maximum: int) -> int: try: parsed = int(value) except (TypeError, ValueError): return default return max(minimum, min(parsed, maximum)) def _safe_bool(value: Any, default: bool) -> bool: if isinstance(value, bool): return value if isinstance(value, (int, float)): return value != 0 if isinstance(value, str): normalized = value.strip().lower() if normalized in { "true", "1", "yes", "on", "enable", "enabled", "开启", "打开", }: return True if normalized in { "false", "0", "no", "off", "disable", "disabled", "关闭", }: return False return default def normalize_point_rules(raw: dict[str, Any] | None) -> dict[str, Any]: rules = dict(DEFAULT_POINT_RULES) if raw: rules.update({key: value for key, value in raw.items() if key in rules}) for key in ( "enabled", "checkin_enabled", "checkin_button_enabled", "activity_enabled", "upvote_enabled", ): rules[key] = _safe_bool(rules[key], DEFAULT_POINT_RULES[key]) timezone_name = str(rules.get("timezone") or "Asia/Shanghai") try: ZoneInfo(timezone_name) except ZoneInfoNotFoundError: timezone_name = "Asia/Shanghai" rules["timezone"] = timezone_name rules["checkin_points"] = _safe_int(rules["checkin_points"], 10, 0, 100000) rules["activity_points"] = _safe_int(rules["activity_points"], 1, 0, 100000) rules["activity_cooldown_seconds"] = _safe_int( rules["activity_cooldown_seconds"], 300, 10, 86400 ) rules["activity_daily_cap"] = _safe_int( rules["activity_daily_cap"], 10, 0, 1000000 ) rules["upvote_points"] = _safe_int(rules["upvote_points"], 2, 0, 100000) rules["upvote_pair_cooldown_seconds"] = _safe_int( rules["upvote_pair_cooldown_seconds"], 86400, 60, 2592000 ) rules["upvote_daily_cap"] = _safe_int( rules["upvote_daily_cap"], 10, 0, 1000000 ) return rules async def ensure_point_indexes() -> None: global _indexes_ready if _indexes_ready: return async with _index_lock: if _indexes_ready: return await accountsdb.create_index( [("chat_id", ASCENDING), ("user_id", ASCENDING)], unique=True ) await accountsdb.create_index( [("chat_id", ASCENDING), ("balance", DESCENDING), ("user_id", ASCENDING)] ) await transactionsdb.create_index([("idempotency_key", ASCENDING)], unique=True) await transactionsdb.create_index( [("chat_id", ASCENDING), ("user_id", ASCENDING), ("created_at", DESCENDING)] ) await transactionsdb.create_index( [("chat_id", ASCENDING), ("source", ASCENDING), ("created_at", DESCENDING)] ) await rulesdb.create_index([("chat_id", ASCENDING)], unique=True) await rate_limitsdb.create_index([("rate_key", ASCENDING)], unique=True) await rate_limitsdb.create_index("expires_at", expireAfterSeconds=0) _indexes_ready = True async def get_point_rules(chat_id: int) -> dict[str, Any]: await ensure_point_indexes() doc = await rulesdb.find_one({"chat_id": chat_id}) return normalize_point_rules(doc) async def set_point_rules(chat_id: int, values: dict[str, Any]) -> dict[str, Any]: await ensure_point_indexes() current = await get_point_rules(chat_id) current.update({key: value for key, value in values.items() if key in DEFAULT_POINT_RULES}) normalized = normalize_point_rules(current) now = utc_now() await rulesdb.update_one( {"chat_id": chat_id}, {"$set": {**normalized, "updated_at": now}, "$setOnInsert": {"created_at": now}}, upsert=True, ) return normalized async def list_checkin_keyboard_migrations( keyboard_version: int, ) -> list[int]: await ensure_point_indexes() cursor = rulesdb.find( { "enabled": True, "checkin_enabled": True, "checkin_button_enabled": True, "$or": [ {"checkin_keyboard_version": {"$exists": False}}, {"checkin_keyboard_version": {"$lt": int(keyboard_version)}}, ], }, {"chat_id": 1}, ) documents = await cursor.to_list(length=10000) return [int(document["chat_id"]) for document in documents] async def mark_checkin_keyboard_version( chat_id: int, keyboard_version: int, ) -> None: await ensure_point_indexes() await rulesdb.update_one( {"chat_id": int(chat_id)}, { "$set": { "checkin_keyboard_version": int(keyboard_version), "updated_at": utc_now(), } }, ) async def _reconcile_account_locked(chat_id: int, user_id: int) -> dict[str, Any]: cached = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id}) pipeline = [ {"$match": {"chat_id": chat_id, "user_id": user_id}}, { "$group": { "_id": None, "balance": {"$sum": "$delta"}, "lifetime_earned": { "$sum": {"$cond": [{"$gt": ["$delta", 0]}, "$delta", 0]} }, "lifetime_spent": { "$sum": { "$cond": [ {"$lt": ["$delta", 0]}, {"$multiply": ["$delta", -1]}, 0, ] } }, } }, ] totals = [doc async for doc in transactionsdb.aggregate(pipeline)] latest = await transactionsdb.find_one( {"chat_id": chat_id, "user_id": user_id}, sort=[("created_at", DESCENDING), ("_id", DESCENDING)], ) total = totals[0] if totals else {} now = utc_now() account = { "chat_id": chat_id, "user_id": user_id, "balance": int(total.get("balance", 0)), "lifetime_earned": int(total.get("lifetime_earned", 0)), "lifetime_spent": int(total.get("lifetime_spent", 0)), "last_transaction_id": latest.get("transaction_id") if latest else None, "updated_at": now, } if latest: account["username"] = ( cached.get("username") if cached and "username" in cached else latest.get("username") ) account["first_name"] = ( cached.get("first_name") if cached and "first_name" in cached else latest.get("first_name") ) account["display_name"] = ( cached.get("display_name") if cached and "display_name" in cached else latest.get("display_name") or latest.get("first_name") ) await accountsdb.update_one( {"chat_id": chat_id, "user_id": user_id}, {"$set": account, "$setOnInsert": {"created_at": now}}, upsert=True, ) return account async def reconcile_account(chat_id: int, user_id: int) -> dict[str, Any]: await ensure_point_indexes() async with _account_locks[(chat_id, user_id)]: return await _reconcile_account_locked(chat_id, user_id) async def get_point_account(chat_id: int, user_id: int) -> dict[str, Any]: await ensure_point_indexes() account = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id}) latest = await transactionsdb.find_one( {"chat_id": chat_id, "user_id": user_id}, projection={"transaction_id": 1}, sort=[("created_at", DESCENDING), ("_id", DESCENDING)], ) if latest and ( not account or account.get("last_transaction_id") != latest.get("transaction_id") ): return await reconcile_account(chat_id, user_id) if account: return account return { "chat_id": chat_id, "user_id": user_id, "balance": 0, "lifetime_earned": 0, "lifetime_spent": 0, "last_transaction_id": None, } async def update_point_account_identity( *, chat_id: int, user_id: int, username: str | None, first_name: str | None, last_name: str | None, display_name: str, ) -> None: await ensure_point_indexes() await accountsdb.update_one( {"chat_id": int(chat_id), "user_id": int(user_id)}, { "$set": { "username": username, "first_name": first_name, "last_name": last_name, "display_name": display_name or first_name, "identity_updated_at": utc_now(), } }, ) async def _write_points_locked( *, chat_id: int, user_id: int, source: str, idempotency_key: str, delta: int | None = None, target_balance: int | None = None, actor_id: int | str | None = None, reason: str = "", reference_id: str | None = None, username: str | None = None, first_name: str | None = None, display_name: str | None = None, ) -> tuple[dict[str, Any], bool]: existing = await transactionsdb.find_one({"idempotency_key": idempotency_key}) if existing: account = await _reconcile_account_locked(chat_id, user_id) return account, False account = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id}) latest = await transactionsdb.find_one( {"chat_id": chat_id, "user_id": user_id}, projection={"transaction_id": 1}, sort=[("created_at", DESCENDING), ("_id", DESCENDING)], ) if not account or ( latest and account.get("last_transaction_id") != latest.get("transaction_id") ): account = await _reconcile_account_locked(chat_id, user_id) balance_before = int(account.get("balance", 0)) actual_delta = ( int(target_balance) - balance_before if target_balance is not None else int(delta or 0) ) balance_after = balance_before + actual_delta if balance_after < 0: raise InsufficientPoints("积分余额不足。") now = utc_now() transaction = { "transaction_id": token_hex(10), "idempotency_key": idempotency_key, "chat_id": chat_id, "user_id": user_id, "username": username, "first_name": first_name, "display_name": display_name or first_name, "delta": actual_delta, "balance_before": balance_before, "balance_after": balance_after, "source": source, "actor_id": actor_id, "reason": reason.strip(), "reference_id": reference_id, "created_at": now, } try: await transactionsdb.insert_one(transaction) except DuplicateKeyError: account = await _reconcile_account_locked(chat_id, user_id) return account, False update = { "$set": { "balance": balance_after, "last_transaction_id": transaction["transaction_id"], "username": username or account.get("username"), "first_name": first_name or account.get("first_name"), "display_name": display_name or first_name or account.get("display_name"), "updated_at": now, }, "$setOnInsert": {"created_at": now}, "$inc": { "lifetime_earned": max(actual_delta, 0), "lifetime_spent": max(-actual_delta, 0), }, } await accountsdb.update_one( {"chat_id": chat_id, "user_id": user_id}, update, upsert=True ) account = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id}) return account or {"chat_id": chat_id, "user_id": user_id, "balance": balance_after}, True async def adjust_points( *, chat_id: int, user_id: int, delta: int, source: str, idempotency_key: str, actor_id: int | str | None = None, reason: str = "", reference_id: str | None = None, username: str | None = None, first_name: str | None = None, display_name: str | None = None, ) -> tuple[dict[str, Any], bool]: await ensure_point_indexes() if source not in POINT_SOURCES: raise PointsError(f"不支持的积分来源:{source}") if not idempotency_key.strip(): raise PointsError("缺少请求幂等键。") delta = int(delta) if delta == 0: account = await get_point_account(chat_id, user_id) return account, False async with _account_locks[(chat_id, user_id)]: return await _write_points_locked( chat_id=chat_id, user_id=user_id, delta=delta, source=source, idempotency_key=idempotency_key, actor_id=actor_id, reason=reason, reference_id=reference_id, username=username, first_name=first_name, display_name=display_name, ) async def set_points( *, chat_id: int, user_id: int, balance: int, actor_id: int | str, reason: str, idempotency_key: str, username: str | None = None, first_name: str | None = None, display_name: str | None = None, ) -> tuple[dict[str, Any], bool]: await ensure_point_indexes() if balance < 0: raise PointsError("积分余额不能小于 0。") if not idempotency_key.strip(): raise PointsError("缺少请求幂等键。") async with _account_locks[(chat_id, user_id)]: return await _write_points_locked( chat_id=chat_id, user_id=user_id, target_balance=balance, source=SOURCE_ADMIN, idempotency_key=idempotency_key, actor_id=actor_id, reason=reason, username=username, first_name=first_name, display_name=display_name, ) async def get_point_transaction_by_key(idempotency_key: str) -> dict[str, Any] | None: await ensure_point_indexes() return await transactionsdb.find_one({"idempotency_key": idempotency_key}) async def list_point_accounts( *, chat_id: int | None = None, query: str = "", page: int = 1, page_size: int = 20, ) -> tuple[list[dict[str, Any]], int]: await ensure_point_indexes() filters: dict[str, Any] = {} if chat_id is not None: filters["chat_id"] = chat_id query = query.strip() if query: if query.lstrip("-").isdigit(): filters["user_id"] = int(query) else: value = re.escape(query.lstrip("@")) filters["$or"] = [ {"username": {"$regex": value, "$options": "i"}}, {"first_name": {"$regex": value, "$options": "i"}}, {"display_name": {"$regex": value, "$options": "i"}}, ] page = max(1, page) page_size = max(1, min(page_size, 100)) total = await accountsdb.count_documents(filters) cursor = ( accountsdb.find(filters) .sort([("balance", DESCENDING), ("user_id", ASCENDING)]) .skip((page - 1) * page_size) .limit(page_size) ) return [doc async for doc in cursor], total async def list_point_transactions( *, chat_id: int | None = None, user_id: int | None = None, source: str | None = None, created_from: datetime | None = None, created_to: datetime | None = None, page: int = 1, page_size: int = 20, ) -> tuple[list[dict[str, Any]], int]: await ensure_point_indexes() filters: dict[str, Any] = {} if chat_id is not None: filters["chat_id"] = chat_id if user_id is not None: filters["user_id"] = user_id if source: if source not in POINT_SOURCES: return [], 0 filters["source"] = source if created_from or created_to: created_filter: dict[str, datetime] = {} if created_from: created_filter["$gte"] = _as_utc(created_from) if created_to: created_filter["$lte"] = _as_utc(created_to) filters["created_at"] = created_filter page = max(1, page) page_size = max(1, min(page_size, 100)) total = await transactionsdb.count_documents(filters) cursor = ( transactionsdb.find(filters) .sort([("created_at", DESCENDING), ("_id", DESCENDING)]) .skip((page - 1) * page_size) .limit(page_size) ) return [doc async for doc in cursor], total async def reconcile_all_accounts( *, chat_id: int | None = None, limit: int = 1000 ) -> dict[str, int]: """Rebuild cached accounts from the immutable ledger.""" await ensure_point_indexes() match: dict[str, Any] = {} if chat_id is not None: match["chat_id"] = int(chat_id) pipeline: list[dict[str, Any]] = [] if match: pipeline.append({"$match": match}) pipeline.extend( [ {"$group": {"_id": {"chat_id": "$chat_id", "user_id": "$user_id"}}}, {"$limit": max(1, min(int(limit), 10000))}, ] ) repaired = 0 async for item in transactionsdb.aggregate(pipeline): await reconcile_account(int(item["_id"]["chat_id"]), int(item["_id"]["user_id"])) repaired += 1 return {"repaired": repaired} async def award_checkin( *, chat_id: int, user_id: int, username: str | None, first_name: str | None, display_name: str | None = None, now: datetime | None = None, ) -> tuple[dict[str, Any], bool]: rules = await get_point_rules(chat_id) if not rules["enabled"] or not rules["checkin_enabled"]: raise PointsError("本群尚未开启积分签到。") now = _as_utc(now or utc_now()) day = _local_day(now, rules["timezone"]) return await adjust_points( chat_id=chat_id, user_id=user_id, delta=rules["checkin_points"], source=SOURCE_CHECKIN, idempotency_key=f"checkin:{chat_id}:{user_id}:{day}", reason=f"每日签到({day})", username=username, first_name=first_name, display_name=display_name, ) async def award_activity( *, chat_id: int, user_id: int, message_id: int, content: str, username: str | None, first_name: str | None, display_name: str | None = None, now: datetime | None = None, ) -> tuple[dict[str, Any] | None, bool]: rules = await get_point_rules(chat_id) if not rules["enabled"] or not rules["activity_enabled"]: return None, False points = int(rules["activity_points"]) cap = int(rules["activity_daily_cap"]) if points <= 0 or cap <= 0: return None, False now = _as_utc(now or utc_now()) day = _local_day(now, rules["timezone"]) rate_key = f"activity:{chat_id}:{user_id}:{day}" content_hash = hashlib.sha256(content.strip().lower().encode("utf-8")).hexdigest() async with _rate_locks[rate_key]: state = await rate_limitsdb.find_one({"rate_key": rate_key}) or {} last_awarded = state.get("last_awarded_at") if last_awarded and ( now - _as_utc(last_awarded) ).total_seconds() < rules["activity_cooldown_seconds"]: return None, False if int(state.get("awarded_points", 0)) >= cap: return None, False recent_hashes = [ item for item in state.get("recent_hashes", []) if now - _as_utc(item["created_at"]) < timedelta(hours=1) ] if any(item["hash"] == content_hash for item in recent_hashes): return None, False award = min(points, cap - int(state.get("awarded_points", 0))) account, created = await adjust_points( chat_id=chat_id, user_id=user_id, delta=award, source=SOURCE_ACTIVITY, idempotency_key=f"activity:{chat_id}:{message_id}:{user_id}", reference_id=str(message_id), reason="活跃消息奖励", username=username, first_name=first_name, display_name=display_name, ) if created: recent_hashes.append({"hash": content_hash, "created_at": now}) await rate_limitsdb.update_one( {"rate_key": rate_key}, { "$set": { "chat_id": chat_id, "user_id": user_id, "rule": SOURCE_ACTIVITY, "last_awarded_at": now, "recent_hashes": recent_hashes[-20:], "expires_at": now + timedelta(days=3), }, "$inc": {"awarded_points": award}, "$setOnInsert": {"created_at": now}, }, upsert=True, ) return account, created async def award_upvote( *, chat_id: int, voter_id: int, target_id: int, message_id: int, username: str | None, first_name: str | None, display_name: str | None = None, now: datetime | None = None, ) -> tuple[dict[str, Any] | None, bool]: if voter_id == target_id: return None, False rules = await get_point_rules(chat_id) if not rules["enabled"] or not rules["upvote_enabled"]: return None, False points = int(rules["upvote_points"]) cap = int(rules["upvote_daily_cap"]) if points <= 0 or cap <= 0: return None, False now = _as_utc(now or utc_now()) day = _local_day(now, rules["timezone"]) daily_key = f"upvote-target:{chat_id}:{target_id}:{day}" pair_key = f"upvote-pair:{chat_id}:{voter_id}:{target_id}" async with _rate_locks[pair_key]: async with _rate_locks[daily_key]: pair = await rate_limitsdb.find_one({"rate_key": pair_key}) or {} last_awarded = pair.get("last_awarded_at") if last_awarded and ( now - _as_utc(last_awarded) ).total_seconds() < rules["upvote_pair_cooldown_seconds"]: return None, False daily = await rate_limitsdb.find_one({"rate_key": daily_key}) or {} already = int(daily.get("awarded_points", 0)) if already >= cap: return None, False award = min(points, cap - already) account, created = await adjust_points( chat_id=chat_id, user_id=target_id, delta=award, source=SOURCE_UPVOTE, idempotency_key=f"upvote:{chat_id}:{message_id}:{voter_id}:{target_id}", actor_id=voter_id, reference_id=str(message_id), reason="收到有效点赞", username=username, first_name=first_name, display_name=display_name, ) if created: expires_at = now + timedelta(days=35) await rate_limitsdb.update_one( {"rate_key": pair_key}, { "$set": { "chat_id": chat_id, "user_id": target_id, "voter_id": voter_id, "rule": "upvote_pair", "last_awarded_at": now, "expires_at": expires_at, }, "$setOnInsert": {"created_at": now}, }, upsert=True, ) await rate_limitsdb.update_one( {"rate_key": daily_key}, { "$set": { "chat_id": chat_id, "user_id": target_id, "rule": SOURCE_UPVOTE, "last_awarded_at": now, "expires_at": now + timedelta(days=3), }, "$inc": {"awarded_points": award}, "$setOnInsert": {"created_at": now}, }, upsert=True, ) return account, created