Tolan/backend/services/energy_stone_service.py

"""
能量石发放与事件记录服务
"""
import uuid
import logging
import time
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional, Tuple

from sqlalchemy import select, func
from sqlalchemy.exc import IntegrityError

from database import db_manager
from models_db import MissionEnergyStone, ConversationEnergyStone, MissionType
from config import settings

logger = logging.getLogger(__name__)

StoneType = Literal["mission", "conversation"]


class EnergyStoneService:
    """能量石生成/状态同步"""

    def __init__(self):
        self.db = db_manager
        self.ttl_seconds = settings.ENERGY_STONE_TTL_SECONDS
        self.conversation_interval_ms = settings.CONVERSATION_STONE_INTERVAL_SECONDS * 1000
        self._mission_type_cache: Dict[str, Tuple[float, Optional[MissionType]]] = {}
        self._mission_type_cache_ttl = max(5, settings.MISSION_TYPE_CACHE_SECONDS)

    async def set_db_manager(self, manager):
        """延迟注入数据库管理器"""
        self.db = manager

    async def create_mission_stone_if_needed(
        self,
        user_id: str,
        request_id: Optional[str],
        mission_type: str = "first_dialog",
        appear_reason: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Optional[Dict[str, Any]]:
        """首次对话奖励石，单用户仅发放一次"""
        if not user_id:
            return None

        try:
            async with self.db.get_session() as session:
                mission_key = mission_type or "first_dialog"
                reason = appear_reason or mission_key
                mission_config: Optional[MissionType] = None
                is_repeatable = False

                async with session.begin():
                    mission_config = await self._get_mission_type_config(session, mission_key)

                    if mission_config is not None and mission_config.is_active is False:
                        logger.debug(
                            "任务类型已禁用，跳过发放: user=%s mission=%s", user_id, mission_key
                        )
                        return None

                    is_repeatable = bool(mission_config.is_repeatable) if mission_config else False

                    result = await session.execute(
                        select(MissionEnergyStone.stone_id)
                        .where(
                            MissionEnergyStone.user_id == user_id,
                            MissionEnergyStone.mission_type == mission_key,
                        )
                        .with_for_update()
                    )
                    existing = result.scalars().first()
                    if existing and not is_repeatable:
                        return None

                    stone = MissionEnergyStone(
                        stone_id=str(uuid.uuid4()),
                        user_id=user_id,
                        request_id=request_id,
                        appear_reason=reason,
                        mission_type=mission_key,
                        metadata=metadata or {},
                    )
                    session.add(stone)

                await session.refresh(stone)

                spawn_meta = {
                    "appear_reason": reason,
                    "mission_type": mission_key,
                    "is_repeatable": is_repeatable,
                }
                if mission_config:
                    spawn_meta["mission_display_name"] = mission_config.display_name

                return self._to_spawn_payload(
                    stone=stone,
                    stone_type="mission",
                    extra=spawn_meta,
                )
        except Exception as exc:
            logger.warning("创建任务能量石失败: user=%s error=%s", user_id, exc)
            return None

    async def create_conversation_stones_if_needed(
        self,
        user_id: str,
        request_id: Optional[str],
        total_duration_ms: int,
        duration_increment_ms: int,
    ) -> List[Dict[str, Any]]:
        """按累计说话时长发放对话能量石"""
        if not user_id or self.conversation_interval_ms <= 0:
            return []

        if total_duration_ms < self.conversation_interval_ms:
            return []

        target_index = total_duration_ms // self.conversation_interval_ms
        if target_index <= 0:
            return []

        stones: List[Dict[str, Any]] = []
        try:
            async with self.db.get_session() as session:
                result = await session.execute(
                    select(func.max(ConversationEnergyStone.threshold_index)).where(
                        ConversationEnergyStone.user_id == user_id
                    )
                )
                last_index = result.scalar() or 0
                next_index = last_index + 1

                if next_index > target_index:
                    return []

                created_instances: List[ConversationEnergyStone] = []
                for idx in range(next_index, target_index + 1):
                    instance = ConversationEnergyStone(
                        stone_id=str(uuid.uuid4()),
                        user_id=user_id,
                        request_id=request_id,
                        threshold_index=idx,
                        total_duration_ms=total_duration_ms,
                        duration_increment_ms=duration_increment_ms,
                        metadata={
                            "target_total_ms": total_duration_ms,
                            "interval_ms": self.conversation_interval_ms,
                        },
                    )
                    session.add(instance)
                    created_instances.append(instance)

                if not created_instances:
                    return []

                await session.commit()
                for instance in created_instances:
                    await session.refresh(instance)
                    stones.append(
                        self._to_spawn_payload(
                            stone=instance,
                            stone_type="conversation",
                            extra={
                                "threshold_index": instance.threshold_index,
                                "total_duration_ms": instance.total_duration_ms,
                                "duration_increment_ms": instance.duration_increment_ms,
                            },
                        )
                    )
            return stones
        except IntegrityError:
            # 并发情况下可能已创建，忽略
            logger.debug("对话能量石已存在，忽略重复创建: user=%s target=%s", user_id, target_index)
            return []
        except Exception as exc:
            logger.warning("创建对话能量石失败: user=%s error=%s", user_id, exc)
            return []

    async def mark_event(
        self,
        stone_id: str,
        event: Literal["click", "timeout"],
        metadata: Optional[Dict[str, Any]] = None,
    ) -> bool:
        """记录点击或自动消失"""
        if not stone_id:
            return False

        try:
            async with self.db.get_session() as session:
                stone = await self._find_stone(session, stone_id)
                if not stone:
                    return False

                if metadata:
                    merged_meta = stone.metadata or {}
                    merged_meta.update(metadata)
                    stone.metadata = merged_meta

                now = datetime.utcnow()
                if event == "click":
                    stone.clicked_at = now
                    stone.auto_removed = False
                else:
                    stone.auto_removed = True
                    if not stone.clicked_at:
                        stone.clicked_at = now

                await session.commit()
                return True
        except Exception as exc:
            logger.warning("记录能量石事件失败: stone=%s event=%s error=%s", stone_id, event, exc)
            return False

    async def _find_stone(self, session, stone_id: str):
        result = await session.execute(
            select(MissionEnergyStone).where(MissionEnergyStone.stone_id == stone_id)
        )
        stone = result.scalars().first()
        if stone:
            return stone

        result = await session.execute(
            select(ConversationEnergyStone).where(ConversationEnergyStone.stone_id == stone_id)
        )
        return result.scalars().first()

    def _to_spawn_payload(
        self,
        stone,
        stone_type: StoneType,
        extra: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        appear_at = stone.appear_at or datetime.utcnow()
        payload = {
            "stone_id": stone.stone_id,
            "stone_type": stone_type,
            "ttl_seconds": self.ttl_seconds,
            "appear_at": int(appear_at.timestamp() * 1000),
            "metadata": extra or {},
        }
        return payload

    async def get_stone_counts(self, user_id: str) -> Dict[str, int]:
        """获取用户的能量石数量"""
        if not user_id:
            return {"mission_count": 0, "conversation_count": 0}

        try:
            async with self.db.get_session() as session:
                # 获取任务能量石数量
                mission_result = await session.execute(
                    select(func.count(MissionEnergyStone.stone_id)).where(
                        MissionEnergyStone.user_id == user_id
                    )
                )
                mission_count = mission_result.scalar() or 0

                # 获取对话能量石数量
                conversation_result = await session.execute(
                    select(func.count(ConversationEnergyStone.stone_id)).where(
                        ConversationEnergyStone.user_id == user_id
                    )
                )
                conversation_count = conversation_result.scalar() or 0

                return {
                    "mission_count": mission_count,
                    "conversation_count": conversation_count,
                }
        except Exception as exc:
            logger.warning("获取能量石数量失败: user=%s error=%s", user_id, exc)
            return {"mission_count": 0, "conversation_count": 0}

    async def _get_mission_type_config(
        self, session, mission_type_key: str
    ) -> Optional[MissionType]:
        """获取任务类型配置，带简单缓存"""
        now = time.monotonic()
        cached = self._mission_type_cache.get(mission_type_key)
        if cached:
            cached_ts, cached_value = cached
            if now - cached_ts < self._mission_type_cache_ttl:
                return cached_value

        result = await session.execute(
            select(MissionType).where(
                MissionType.type_key == mission_type_key,
                MissionType.is_active.is_(True),
            )
        )
        mission_type = result.scalars().first()
        self._mission_type_cache[mission_type_key] = (now, mission_type)
        return mission_type


energy_stone_service = EnergyStoneService()