"""Token 预算追踪器 — 从 ai_run_logs 聚合日/月 token 消耗。

每次 AI 调用前检查预算，超限时拒绝请求。
日预算默认 100,000 tokens，月预算默认 2,000,000 tokens。

聚合数据通过构造函数注入的 callable 获取（解耦 AIRunLogService），
callable 签名：() -> int，分别返回当日/当月已消耗 token 数。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Protocol


class UsageProvider(Protocol):
    """Token 用量数据提供者协议。"""

    def get_daily_usage(self) -> int:
        """返回当日已消耗 token 数。"""
        ...

    def get_monthly_usage(self) -> int:
        """返回当月已消耗 token 数。"""
        ...


@dataclass
class BudgetStatus:
    """预算检查结果。"""

    allowed: bool
    daily_used: int
    monthly_used: int
    reason: str | None = None  # "daily_exceeded" / "monthly_exceeded" / None


class BudgetTracker:
    """Token 预算追踪器，从 ai_run_logs 聚合。

    支持两种注入方式：
    1. 传入 UsageProvider 实例（如 AIRunLogService）
    2. 传入两个 callable：get_daily_usage / get_monthly_usage
    """

    def __init__(
        self,
        daily_limit: int = 100_000,
        monthly_limit: int = 2_000_000,
        *,
        get_daily_usage: Callable[[], int] | None = None,
        get_monthly_usage: Callable[[], int] | None = None,
        usage_provider: UsageProvider | None = None,
    ) -> None:
        self.daily_limit = daily_limit
        self.monthly_limit = monthly_limit

        # 优先使用 usage_provider，其次使用独立 callable
        if usage_provider is not None:
            self._get_daily_usage = usage_provider.get_daily_usage
            self._get_monthly_usage = usage_provider.get_monthly_usage
        elif get_daily_usage is not None and get_monthly_usage is not None:
            self._get_daily_usage = get_daily_usage
            self._get_monthly_usage = get_monthly_usage
        else:
            raise ValueError(
                "必须提供 usage_provider 或同时提供 "
                "get_daily_usage 和 get_monthly_usage callable"
            )

    def check_budget(self) -> BudgetStatus:
        """检查当前预算状态。

        先检查日预算，再检查月预算。
        任一超限即返回 allowed=False 并附带原因。
        """
        daily_used = self._get_daily_usage()
        monthly_used = self._get_monthly_usage()

        if daily_used >= self.daily_limit:
            return BudgetStatus(
                allowed=False,
                daily_used=daily_used,
                monthly_used=monthly_used,
                reason="daily_exceeded",
            )

        if monthly_used >= self.monthly_limit:
            return BudgetStatus(
                allowed=False,
                daily_used=daily_used,
                monthly_used=monthly_used,
                reason="monthly_exceeded",
            )

        return BudgetStatus(
            allowed=True,
            daily_used=daily_used,
            monthly_used=monthly_used,
            reason=None,
        )