zhenxun_bot/zhenxun/services/auth_snapshot/service.py

"""
快照服务

提供权限快照的获取、缓存、失效等功能
"""

import asyncio
from typing import ClassVar

from zhenxun.services.cache import CacheRoot, cache_config
from zhenxun.services.cache.cache_containers import CacheDict
from zhenxun.services.cache.config import CacheMode
from zhenxun.services.log import logger
from zhenxun.utils.enum import CacheType

from .builder import SnapshotBuilder
from .models import AuthSnapshot, PluginSnapshot

LOG_COMMAND = "auth_snapshot"

# 内存缓存名称（CacheType 已提供 Redis 前缀，此处仅用于内存缓存标识）
AUTH_MEMORY_CACHE_NAME = "AUTH_MEMORY"
PLUGIN_MEMORY_CACHE_NAME = "PLUGIN_MEMORY"

# 内存缓存TTL配置
AUTH_MEMORY_TTL = 10  # 权限快照内存缓存TTL（秒）
AUTH_REDIS_TTL = 60  # 权限快照Redis缓存TTL（秒）
PLUGIN_MEMORY_TTL = 30  # 插件快照内存缓存TTL（秒）
PLUGIN_REDIS_TTL = 300  # 插件快照Redis缓存TTL（秒）

# 并发控制配置
MAX_CONCURRENT_BUILDS = 15  # 最大同时构建数量（防止 DB 过载）
BUILD_QUEUE_TIMEOUT = 5.0  # 等待构建队列的超时时间（秒）


class AuthSnapshotService:
    """权限快照服务

    提供权限快照的获取、缓存和失效管理
    """

    # 本地内存缓存（使用 CacheDict，自动处理过期）
    _memory_cache: ClassVar[CacheDict[AuthSnapshot] | None] = None

    # 正在构建中的快照（防止并发重复构建）
    _building: ClassVar[dict[str, asyncio.Future]] = {}

    # per-key 锁（保护 _building 的检查和设置，防止竞态条件）
    _build_locks: ClassVar[dict[str, asyncio.Lock]] = {}

    # 全局构建并发限制（防止大量不同 key 同时构建导致 DB 过载）
    _build_semaphore: ClassVar[asyncio.Semaphore | None] = None

    @classmethod
    def _get_build_semaphore(cls) -> asyncio.Semaphore:
        """获取构建信号量（懒加载）"""
        if cls._build_semaphore is None:
            cls._build_semaphore = asyncio.Semaphore(MAX_CONCURRENT_BUILDS)
        return cls._build_semaphore

    @classmethod
    def _get_memory_cache(cls) -> CacheDict[AuthSnapshot]:
        """获取内存缓存实例（懒加载）"""
        if cls._memory_cache is None:
            cls._memory_cache = CacheRoot.cache_dict(
                AUTH_MEMORY_CACHE_NAME,
                expire=AUTH_MEMORY_TTL,
                value_type=AuthSnapshot,
            )
        return cls._memory_cache

    @classmethod
    def _build_cache_key(cls, user_id: str, group_id: str | None, bot_id: str) -> str:
        """构建缓存键（CacheType 已提供前缀，此处只需业务标识）"""
        group_part = group_id or "PRIVATE"
        return f"{user_id}:{group_part}:{bot_id}"

    @classmethod
    async def get_snapshot(
        cls,
        user_id: str,
        group_id: str | None,
        bot_id: str,
        force_refresh: bool = False,
    ) -> AuthSnapshot:
        """获取权限快照

        优先从缓存获取，缓存未命中时构建新快照

        参数:
            user_id: 用户ID
            group_id: 群组ID（可为None）
            bot_id: Bot ID
            force_refresh: 是否强制刷新

        返回:
            AuthSnapshot: 权限快照
        """
        cache_key = cls._build_cache_key(user_id, group_id, bot_id)

        memory_cache = cls._get_memory_cache()

        # 1. 尝试从内存缓存获取（最快路径）
        if not force_refresh:
            if snapshot := memory_cache.get(cache_key):
                return snapshot

        # 2. 尝试从Redis获取
        if not force_refresh and cache_config.cache_mode != CacheMode.NONE:
            try:
                cached = await CacheRoot.get(CacheType.AUTH_SNAPSHOT, cache_key)
                if cached and isinstance(cached, dict):
                    snapshot = AuthSnapshot.model_validate(cached)
                    if not snapshot.is_expired(AUTH_REDIS_TTL):
                        memory_cache.set(cache_key, snapshot)
                        return snapshot
            except Exception as e:
                logger.debug(f"从Redis获取快照失败: {cache_key}", LOG_COMMAND, e=e)

        # 3. 获取或创建 per-key 锁（使用 setdefault 保证原子性）
        lock = cls._build_locks.setdefault(cache_key, asyncio.Lock())

        # 4. 先尝试快速路径：检查是否有其他协程正在构建
        if cache_key in cls._building:
            try:
                return await cls._building[cache_key]
            except Exception:
                pass

        # 5. 获取信号量（控制总并发数，不在锁内等待）
        semaphore = cls._get_build_semaphore()
        try:
            await asyncio.wait_for(semaphore.acquire(), timeout=BUILD_QUEUE_TIMEOUT)
        except asyncio.TimeoutError:
            logger.warning(f"获取信号量超时，使用默认快照: {cache_key}", LOG_COMMAND)
            return AuthSnapshot(user_id=user_id, group_id=group_id, bot_id=bot_id)

        need_build = False
        future: asyncio.Future[AuthSnapshot] | None = None

        try:
            # 6. 获取 per-key 锁，只保护 _building 的检查和设置
            async with lock:
                # 再次检查缓存
                if snapshot := memory_cache.get(cache_key):
                    return snapshot

                # 检查是否有其他协程正在构建
                if cache_key in cls._building:
                    future = cls._building[cache_key]
                else:
                    # 创建 future 并设置到 _building（在锁内）
                    loop = asyncio.get_running_loop()
                    future = loop.create_future()
                    cls._building[cache_key] = future
                    need_build = True

            # 7. 锁外执行（构建或等待）
            if need_build:
                return await cls._do_build_with_future(
                    user_id, group_id, bot_id, cache_key, future
                )
            else:
                # 等待其他协程的构建结果
                return await future  # type: ignore
        finally:
            semaphore.release()

    @classmethod
    async def _do_build_with_future(
        cls,
        user_id: str,
        group_id: str | None,
        bot_id: str,
        cache_key: str,
        future: asyncio.Future[AuthSnapshot],
    ) -> AuthSnapshot:
        """执行快照构建（future 已在锁内设置到 _building）"""
        try:
            # 构建快照
            snapshot = await SnapshotBuilder.build_auth_snapshot(
                user_id, group_id, bot_id
            )

            # 存入Redis缓存（异步，不阻塞）
            if cache_config.cache_mode != CacheMode.NONE:
                asyncio.create_task(  # noqa: RUF006
                    cls._cache_to_redis(cache_key, snapshot)
                )

            # 存入内存缓存
            cls._get_memory_cache().set(cache_key, snapshot)

            future.set_result(snapshot)
            return snapshot

        except Exception as e:
            future.set_exception(e)
            raise
        finally:
            cls._building.pop(cache_key, None)

    @classmethod
    async def _cache_to_redis(cls, cache_key: str, snapshot: AuthSnapshot):
        """异步存入Redis"""
        try:
            await CacheRoot.set(
                CacheType.AUTH_SNAPSHOT,
                cache_key,
                snapshot.model_dump(),
                expire=AUTH_REDIS_TTL,
            )
        except Exception as e:
            logger.debug(f"缓存权限快照到Redis失败: {cache_key}", LOG_COMMAND, e=e)

    @classmethod
    async def invalidate_user(cls, user_id: str):
        """失效用户相关的所有快照

        参数:
            user_id: 用户ID
        """
        # 清理内存缓存（遍历 CacheDict 的 keys）
        memory_cache = cls._get_memory_cache()
        keys_to_delete = [k for k in memory_cache.keys() if f":{user_id}:" in k]
        for key in keys_to_delete:
            del memory_cache[key]

        logger.debug(f"已失效用户 {user_id} 的权限快照缓存", LOG_COMMAND)

    @classmethod
    async def invalidate_group(cls, group_id: str):
        """失效群组相关的所有快照

        参数:
            group_id: 群组ID
        """
        # 清理内存缓存
        memory_cache = cls._get_memory_cache()
        keys_to_delete = [k for k in memory_cache.keys() if f":{group_id}:" in k]
        for key in keys_to_delete:
            del memory_cache[key]

        logger.debug(f"已失效群组 {group_id} 的权限快照缓存", LOG_COMMAND)

    @classmethod
    async def invalidate_bot(cls, bot_id: str):
        """失效Bot相关的所有快照

        参数:
            bot_id: Bot ID
        """
        # 清理内存缓存
        memory_cache = cls._get_memory_cache()
        keys_to_delete = [k for k in memory_cache.keys() if k.endswith(f":{bot_id}")]
        for key in keys_to_delete:
            del memory_cache[key]

        logger.debug(f"已失效Bot {bot_id} 的权限快照缓存", LOG_COMMAND)

    @classmethod
    def clear_all_cache(cls):
        """清空所有缓存"""
        if cls._memory_cache:
            cls._memory_cache.clear()
        cls._building.clear()
        cls._build_locks.clear()
        logger.info("已清空所有权限快照缓存", LOG_COMMAND)


class PluginSnapshotService:
    """插件快照服务

    提供插件信息的获取和缓存，支持本地内存缓存 + Redis 双层缓存
    """

    # 本地内存缓存（使用 CacheDict）
    _memory_cache: ClassVar[CacheDict[PluginSnapshot] | None] = None

    # 正在构建中的快照
    _building: ClassVar[dict[str, asyncio.Future]] = {}

    @classmethod
    def _get_memory_cache(cls) -> CacheDict[PluginSnapshot]:
        """获取内存缓存实例（懒加载）"""
        if cls._memory_cache is None:
            cls._memory_cache = CacheRoot.cache_dict(
                PLUGIN_MEMORY_CACHE_NAME,
                expire=PLUGIN_MEMORY_TTL,
                value_type=PluginSnapshot,
            )
        return cls._memory_cache

    @classmethod
    def _build_cache_key(cls, module: str) -> str:
        """构建缓存键（CacheType 已提供前缀，此处只需模块名）"""
        return module

    @classmethod
    async def get_plugin(
        cls, module: str, force_refresh: bool = False
    ) -> PluginSnapshot | None:
        """获取插件快照

        参数:
            module: 插件模块名
            force_refresh: 是否强制刷新

        返回:
            PluginSnapshot | None: 插件快照，不存在时返回None
        """
        cache_key = cls._build_cache_key(module)
        memory_cache = cls._get_memory_cache()

        # 1. 尝试从内存缓存获取（最快路径）
        if not force_refresh:
            if snapshot := memory_cache.get(cache_key):
                return snapshot

        # 2. 尝试从Redis获取
        if not force_refresh and cache_config.cache_mode != CacheMode.NONE:
            try:
                cached = await CacheRoot.get(CacheType.PLUGIN_SNAPSHOT, cache_key)
                if cached and isinstance(cached, dict):
                    snapshot = PluginSnapshot.model_validate(cached)
                    if not snapshot.is_expired(PLUGIN_REDIS_TTL):
                        memory_cache.set(cache_key, snapshot)
                        return snapshot
            except Exception as e:
                logger.debug(f"从Redis获取插件快照失败: {module}", LOG_COMMAND, e=e)

        # 3. 检查是否有其他协程正在构建
        if cache_key in cls._building:
            try:
                return await cls._building[cache_key]
            except Exception:
                pass

        # 4. 从数据库构建（插件数量有限，无需信号量）
        return await cls._do_build(module, cache_key)

    @classmethod
    async def _do_build(cls, module: str, cache_key: str) -> PluginSnapshot | None:
        """执行插件快照构建"""
        loop = asyncio.get_running_loop()
        future: asyncio.Future[PluginSnapshot | None] = loop.create_future()
        cls._building[cache_key] = future

        try:
            snapshot = await SnapshotBuilder.build_plugin_snapshot(module)

            if snapshot:
                # 存入Redis缓存（异步）
                if cache_config.cache_mode != CacheMode.NONE:
                    asyncio.create_task(  # noqa: RUF006
                        cls._cache_to_redis(cache_key, snapshot)
                    )

                # 存入内存缓存
                cls._get_memory_cache().set(cache_key, snapshot)

            future.set_result(snapshot)
            return snapshot

        except Exception as e:
            future.set_exception(e)
            raise
        finally:
            cls._building.pop(cache_key, None)

    @classmethod
    async def _cache_to_redis(cls, cache_key: str, snapshot: PluginSnapshot):
        """异步存入Redis"""
        try:
            await CacheRoot.set(
                CacheType.PLUGIN_SNAPSHOT,
                cache_key,
                snapshot.model_dump(),
                expire=PLUGIN_REDIS_TTL,
            )
        except Exception as e:
            logger.debug(f"缓存插件快照到Redis失败: {cache_key}", LOG_COMMAND, e=e)

    @classmethod
    async def invalidate_plugin(cls, module: str):
        """失效指定插件的缓存

        参数:
            module: 插件模块名
        """
        cache_key = cls._build_cache_key(module)

        # 清理内存缓存
        memory_cache = cls._get_memory_cache()
        if cache_key in memory_cache.keys():
            del memory_cache[cache_key]

        # 清理Redis缓存
        if cache_config.cache_mode != CacheMode.NONE:
            try:
                await CacheRoot.delete(CacheType.PLUGIN_SNAPSHOT, cache_key)
            except Exception:
                pass

        logger.debug(f"已失效插件 {module} 的快照缓存", LOG_COMMAND)

    @classmethod
    async def warmup(cls):
        """预热所有插件缓存

        在启动时调用，预加载所有插件信息到缓存
        同时写入内存缓存和 Redis 缓存
        """
        from zhenxun.models.plugin_info import PluginInfo

        memory_cache = cls._get_memory_cache()

        try:
            plugins = await PluginInfo.filter(load_status=True).all()
            count = 0

            for plugin in plugins:
                snapshot = PluginSnapshot(
                    module=plugin.module,
                    name=plugin.name,
                    status=plugin.status,
                    block_type=plugin.block_type,
                    plugin_type=plugin.plugin_type,
                    admin_level=plugin.admin_level or 0,
                    cost_gold=plugin.cost_gold,
                    level=plugin.level,
                    limit_superuser=plugin.limit_superuser,
                    ignore_prompt=plugin.ignore_prompt,
                )

                cache_key = cls._build_cache_key(plugin.module)

                # 存入内存缓存（最快访问路径）
                memory_cache.set(cache_key, snapshot)

                # 同时存入 Redis 缓存（跨进程共享）
                if cache_config.cache_mode != CacheMode.NONE:
                    try:
                        await CacheRoot.set(
                            CacheType.PLUGIN_SNAPSHOT,
                            cache_key,
                            snapshot,
                            expire=PLUGIN_REDIS_TTL,
                        )
                    except Exception:
                        pass  # Redis 写入失败不影响预热

                count += 1

            logger.info(f"已预热 {count} 个插件的快照缓存", LOG_COMMAND)

        except Exception as e:
            logger.error("预热插件缓存失败", LOG_COMMAND, e=e)

    @classmethod
    def clear_all_cache(cls):
        """清空所有缓存"""
        if cls._memory_cache:
            cls._memory_cache.clear()
        cls._building.clear()
        logger.info("已清空所有插件快照缓存", LOG_COMMAND)