"""用量记录器（Redis 持久化）.

按维度记录 token 用量：
  - user:{user_id}:usage:{YYYYMM} → 月度用户用量
  - school:{school_id}:usage:{YYYYMM} → 月度学校用量

全并行模式：Redis 不可用时降级（仅记录日志，不阻断主流程）。
"""

from dataclasses import dataclass
from datetime import UTC, datetime

import structlog
from redis.asyncio import Redis
from redis.exceptions import RedisError

logger = structlog.get_logger()


@dataclass
class UsageRecord:
    """单次用量记录."""

    user_id: str
    school_id: str
    provider: str
    model: str
    operation: str
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    latency_ms: int = 0
    success: bool = True
    degraded: bool = False
    request_id: str = ""


class UsageRecorder:
    """用量记录器.

    将每次 LLM 调用的 token 用量记录到 Redis，供 QuotaEnforcer 查询。
    全并行模式：Redis 不可用时降级，仅记录日志。
    """

    def __init__(self, redis: Redis | None = None) -> None:
        self._redis = redis

    async def record(self, record: UsageRecord) -> None:
        """记录一次用量.

        Args:
            record: 用量记录
        """
        if self._redis is None:
            logger.debug(
                "usage_record_skipped_no_redis",
                user_id=record.user_id,
                operation=record.operation,
                total_tokens=record.total_tokens,
            )
            return

        month_key = self._month_key()
        user_key = f"user:{record.user_id}:usage:{month_key}"
        school_key = f"school:{record.school_id}:usage:{month_key}"

        try:
            # 使用 INCRBY 原子递增 token 用量
            pipe = self._redis.pipeline()
            pipe.incrby(user_key, record.total_tokens)
            pipe.incrby(school_key, record.total_tokens)
            # 设置过期时间（35 天，覆盖月度周期）
            pipe.expire(user_key, 35 * 86400)
            pipe.expire(school_key, 35 * 86400)
            await pipe.execute()

            logger.info(
                "usage_recorded",
                user_id=record.user_id,
                school_id=record.school_id,
                operation=record.operation,
                provider=record.provider,
                model=record.model,
                total_tokens=record.total_tokens,
                success=record.success,
                degraded=record.degraded,
            )
        except RedisError as exc:
            logger.warning(
                "usage_record_failed_redis_error",
                error=str(exc),
                user_id=record.user_id,
            )

    async def get_user_usage(self, user_id: str, month: str = "") -> int:
        """查询用户月度用量.

        Args:
            user_id: 用户 ID
            month: 月份键（空则取当前月）

        Returns:
            已用 token 数（Redis 不可用时返回 0）
        """
        if self._redis is None:
            return 0
        month_key = month or self._month_key()
        key = f"user:{user_id}:usage:{month_key}"
        try:
            value = await self._redis.get(key)
            return int(value) if value else 0
        except RedisError as exc:
            logger.warning("usage_query_failed", error=str(exc))
            return 0

    async def get_school_usage(self, school_id: str, month: str = "") -> int:
        """查询学校月度用量.

        Args:
            school_id: 学校 ID
            month: 月份键（空则取当前月）

        Returns:
            已用 token 数（Redis 不可用时返回 0）
        """
        if self._redis is None:
            return 0
        month_key = month or self._month_key()
        key = f"school:{school_id}:usage:{month_key}"
        try:
            value = await self._redis.get(key)
            return int(value) if value else 0
        except RedisError as exc:
            logger.warning("usage_query_failed", error=str(exc))
            return 0

    @staticmethod
    def _month_key() -> str:
        """当前月份键（YYYYMM 格式）."""
        return datetime.now(UTC).strftime("%Y%m")
