# -*- coding: utf-8 -*-
"""
思考预算计算器

借鉴 zhikuncode-main `llm/ThinkingBudgetCalculator.java`，基于上一轮对话的
三因子动态计算本轮 thinking token 预算：

- Factor 1（工具调用密度）：≥3 次 ×1.5，≥1 次 ×1.2，否则 ×1.0
- Factor 2（输出 token 量级）：>2000 ×1.3，>500 ×1.1，否则 ×1.0
- Factor 3（对话轮次深度）：>10 轮 ×1.2，>5 轮 ×1.1，否则 ×1.0

三因子相乘 × BASE_BUDGET(10000) → clamp 到 [2000, 50000]。

DeepSeek API 使用 `reasoning_effort` 字符串参数（low/medium/high/max），
本模块同时提供 `budget_to_effort()` 把数值预算映射到字符串。
"""
import threading
from dataclasses import dataclass, asdict
from typing import Optional


@dataclass
class ContextMetrics:
    """上一轮对话的上下文度量，用于计算本轮思考预算。

    所有字段在首次调用（无历史）时为 0/0/1。
    """
    tool_calls: int = 0       # 上一轮工具调用次数
    output_tokens: int = 0    # 上一轮输出 token 数
    turn_count: int = 1       # 当前对话轮次（从 1 开始）

    @classmethod
    def initial(cls) -> "ContextMetrics":
        """首轮默认值：所有 factor = 1.0"""
        return cls(tool_calls=0, output_tokens=0, turn_count=1)

    def to_dict(self) -> dict:
        return asdict(self)


class ThinkingBudgetCalculator:
    """三因子自适应思考预算计算器。

    线程安全：纯计算无状态，可全局共享单例。
    """

    BASE_BUDGET: int = 10000
    MIN: int = 2000
    MAX: int = 50000

    # 三因子阶梯阈值
    TOOL_CALL_HIGH = 3       # ≥3 次工具调用触发高系数
    TOOL_CALL_MED = 1        # ≥1 次工具调用触发中系数
    OUTPUT_TOKENS_HIGH = 2000  # >2000 token 触发高系数
    OUTPUT_TOKENS_MED = 500    # >500 token 触发中系数
    TURN_COUNT_HIGH = 10      # >10 轮触发高系数
    TURN_COUNT_MED = 5        # >5 轮触发中系数

    # 三因子系数
    FACTOR_HIGH = 1.5
    FACTOR_MED = 1.2
    FACTOR_LOW = 1.0

    FACTOR_OUTPUT_HIGH = 1.3
    FACTOR_OUTPUT_MED = 1.1
    FACTOR_OUTPUT_LOW = 1.0

    FACTOR_TURN_HIGH = 1.2
    FACTOR_TURN_MED = 1.1
    FACTOR_TURN_LOW = 1.0

    def calculate(self, metrics: ContextMetrics) -> int:
        """计算本轮思考预算。

        Args:
            metrics: 上一轮对话的上下文度量

        Returns:
            思考 token 预算，clamp 到 [MIN, MAX]
        """
        # 负值兜底为 0
        tool_calls = max(0, metrics.tool_calls)
        output_tokens = max(0, metrics.output_tokens)
        turn_count = max(1, metrics.turn_count)

        # Factor 1：工具调用密度
        if tool_calls >= self.TOOL_CALL_HIGH:
            f1 = self.FACTOR_HIGH
        elif tool_calls >= self.TOOL_CALL_MED:
            f1 = self.FACTOR_MED
        else:
            f1 = self.FACTOR_LOW

        # Factor 2：输出 token 量级
        if output_tokens > self.OUTPUT_TOKENS_HIGH:
            f2 = self.FACTOR_OUTPUT_HIGH
        elif output_tokens > self.OUTPUT_TOKENS_MED:
            f2 = self.FACTOR_OUTPUT_MED
        else:
            f2 = self.FACTOR_OUTPUT_LOW

        # Factor 3：对话轮次深度
        if turn_count > self.TURN_COUNT_HIGH:
            f3 = self.FACTOR_TURN_HIGH
        elif turn_count > self.TURN_COUNT_MED:
            f3 = self.FACTOR_TURN_MED
        else:
            f3 = self.FACTOR_TURN_LOW

        budget = int(self.BASE_BUDGET * f1 * f2 * f3)
        return max(self.MIN, min(self.MAX, budget))


# budget → reasoning_effort 字符串映射
# DeepSeek API 支持：low / medium / high / max
_EFFORT_THRESHOLDS = [
    (50000, "max"),     # >=50000 → max
    (25000, "high"),    # >=25000 → high
    (10000, "medium"),  # >=10000 → medium
    (0, "low"),         # <10000 → low
]


def budget_to_effort(budget: int) -> str:
    """把数值思考预算映射到 reasoning_effort 字符串。

    DeepSeek API 的 `reasoning_effort` 参数接受字符串值，
    本函数把数值预算映射到最近的 effort 等级。

    映射规则：
    - 50000+ → "max"
    - 25000-49999 → "high"
    - 10000-24999 → "medium"
    - 0-9999 → "low"

    Args:
        budget: 思考 token 预算

    Returns:
        effort 字符串（"low" / "medium" / "high" / "max"）
    """
    for threshold, label in _EFFORT_THRESHOLDS:
        if budget >= threshold:
            return label
    return "low"  # 兜底


# 模块级单例
_default_calc: Optional[ThinkingBudgetCalculator] = None
_default_lock = threading.Lock()


def get_default_calculator() -> ThinkingBudgetCalculator:
    """获取默认单例 ThinkingBudgetCalculator。"""
    global _default_calc
    if _default_calc is None:
        with _default_lock:
            if _default_calc is None:
                _default_calc = ThinkingBudgetCalculator()
    return _default_calc


def compute_adaptive_effort(
    metrics: ContextMetrics,
    current_effort: str,
    enabled: bool = False,
) -> str:
    """根据上下文度量计算自适应 reasoning_effort。

    若 enabled=False 或 current_effort 为 "max"（用户显式指定最高档，
    不应被自动降级），返回 current_effort 原值。
    否则计算数值预算并映射到 effort 字符串。

    Args:
        metrics: 上一轮对话的上下文度量
        current_effort: 当前配置的 reasoning_effort（如 "high"）
        enabled: 是否启用自适应（来自 config.adaptive_thinking_budget）

    Returns:
        effort 字符串（"low" / "medium" / "high" / "max"）
    """
    if not enabled:
        return current_effort
    # "max" 是用户显式指定的最高档，不应被自动降级
    if current_effort == "max":
        return current_effort
    budget = get_default_calculator().calculate(metrics)
    return budget_to_effort(budget)
