zhenxun_bot/zhenxun/builtin_plugins/hooks/auth_side_effect.py

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass, field
import time
from typing import Any, Protocol

from nonebot_plugin_uninfo import Uninfo

from zhenxun.services.log import logger
from zhenxun.utils.utils import EntityIDs

from .auth.config import LOGGER_COMMAND
from .auth.utils import send_message

AsyncAction = Callable[[], Awaitable[None]]


class SyncReservation(Protocol):
    def commit(self) -> None: ...

    def release(self) -> None: ...


class AsyncReservation(Protocol):
    async def commit(self) -> None: ...

    async def release(self) -> None: ...


ReservationLike = AsyncAction | SyncReservation | AsyncReservation
SideEffectKind = str
SideEffectState = str


@dataclass(slots=True)
class SideEffectReservation:
    kind: SideEffectKind
    reservation: ReservationLike
    amount: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)
    state: SideEffectState = "reserved"
    reserved_at: float = field(default_factory=time.monotonic)
    committed_at: float | None = None
    released_at: float | None = None
    reason: str | None = None

    @property
    def should_auto_unblock(self) -> bool:
        return bool(getattr(self.reservation, "should_auto_unblock", False))


async def _maybe_await(value: Any) -> None:
    if hasattr(value, "__await__"):
        await value


async def _commit_reservation(reservation: ReservationLike) -> None:
    commit = getattr(reservation, "commit", None)
    if callable(commit):
        await _maybe_await(commit())
        return
    if callable(reservation):
        await reservation()


async def _release_reservation(reservation: ReservationLike) -> None:
    release = getattr(reservation, "release", None)
    if callable(release):
        await _maybe_await(release())


@dataclass(slots=True)
class SideEffectCommit:
    """权限链副作用提交器。

    第一阶段只封装既有调用点，不改变扣金币、限流提交、权限提示发送时机。
    """

    session: Uninfo
    module: str
    owner_matcher_id: int | None = None
    limit_entity: EntityIDs | None = None
    _reservations: dict[SideEffectKind, SideEffectReservation] = field(
        default_factory=dict
    )
    committed: bool = False

    @property
    def limit_should_auto_unblock(self) -> bool:
        record = self._reservations.get("limit")
        return bool(record and record.should_auto_unblock)

    @property
    def has_pending(self) -> bool:
        return any(record.state == "reserved" for record in self._reservations.values())

    @property
    def pending_kinds(self) -> tuple[str, ...]:
        return tuple(
            kind
            for kind, record in self._reservations.items()
            if record.state == "reserved"
        )

    def snapshot(self) -> dict[str, Any]:
        return {
            "module": self.module,
            "committed": self.committed,
            "pending": list(self.pending_kinds),
            "reservations": {
                kind: {
                    "state": record.state,
                    "amount": record.amount,
                    "metadata": record.metadata,
                    "reason": record.reason,
                }
                for kind, record in self._reservations.items()
            },
        }

    async def send_permission_tip(
        self,
        message: list | str,
        check_tag: str | None = None,
        *,
        background: bool = False,
        timeout: float | None = None,
    ) -> None:
        try:
            tip_coro = send_message(
                self.session,
                message,
                check_tag,
                background=background,
            )
            if timeout and not background:
                await asyncio.wait_for(tip_coro, timeout=timeout)
            else:
                await tip_coro
        except asyncio.TimeoutError:
            logger.error("发送权限提示超时", LOGGER_COMMAND, session=self.session)

    async def reduce_gold(
        self,
        func: ReservationLike,
    ) -> None:
        await self.reserve_gold(func)
        await self.commit_gold()

    async def reserve(
        self,
        kind: SideEffectKind,
        reservation: ReservationLike,
        *,
        amount: int = 0,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.release(kind, f"replace_{kind}_reservation")
        self._reservations[kind] = SideEffectReservation(
            kind=kind,
            reservation=reservation,
            amount=amount,
            metadata=metadata or {},
        )

    async def commit(self, kind: SideEffectKind) -> None:
        record = self._reservations.get(kind)
        if record is None or record.state != "reserved":
            return
        try:
            await _commit_reservation(record.reservation)
        except Exception:
            record.reason = "commit_failed"
            raise
        record.state = "committed"
        record.committed_at = time.monotonic()

    async def release(
        self,
        kind: SideEffectKind,
        reason: str | None = None,
    ) -> None:
        record = self._reservations.get(kind)
        if record is None or record.state != "reserved":
            return
        try:
            await _release_reservation(record.reservation)
        finally:
            record.state = "released"
            record.released_at = time.monotonic()
            record.reason = reason

    async def reserve_limit(self, reservation: ReservationLike) -> None:
        await self.reserve("limit", reservation)

    async def commit_limit(
        self,
        reservation: ReservationLike | None = None,
    ) -> None:
        if reservation is not None:
            await self.reserve_limit(reservation)
        await self.commit("limit")

    async def release_limit(self, reason: str | None = None) -> None:
        await self.release("limit", reason)

    async def reserve_gold(
        self,
        reservation: ReservationLike,
        *,
        amount: int = 0,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.reserve(
            "gold",
            reservation,
            amount=amount,
            metadata=metadata,
        )

    async def commit_gold(self) -> None:
        await self.commit("gold")

    async def rollback_gold(self, reason: str | None = None) -> None:
        await self.release("gold", reason)

    async def rollback_all(self, reason: str | None = None) -> None:
        for kind in list(self._reservations):
            await self.release(kind, reason)

    async def commit_all(self, *, order: Sequence[str] = ("gold", "limit")) -> None:
        for name in order:
            await self.commit(name)
        self.committed = True