TelegramBotAdmin/wbb/utils/dbpoints.py

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 _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