# -*- coding: utf-8 -*-
"""ThinkingBudgetCalculator 测试 — 三因子自适应思考预算"""
import pytest

from abyss.llm.thinking_budget import (
    ContextMetrics,
    ThinkingBudgetCalculator,
    budget_to_effort,
)


@pytest.fixture
def calc():
    return ThinkingBudgetCalculator()


# -------- 单因子边界测试 --------

def test_initial_metrics_returns_base_budget(calc):
    """首轮默认所有 factor = 1.0，结果等于 BASE_BUDGET"""
    metrics = ContextMetrics.initial()
    budget = calc.calculate(metrics)
    assert budget == ThinkingBudgetCalculator.BASE_BUDGET


def test_no_tool_calls_no_boost(calc):
    """0 次工具调用，f1 = 1.0"""
    metrics = ContextMetrics(tool_calls=0, output_tokens=100, turn_count=1)
    budget = calc.calculate(metrics)
    assert budget == ThinkingBudgetCalculator.BASE_BUDGET  # 10000


def test_one_tool_call_boosts_f1(calc):
    """1 次工具调用，f1 = 1.2"""
    metrics = ContextMetrics(tool_calls=1, output_tokens=100, turn_count=1)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.2)


def test_three_tool_calls_boosts_f1_more(calc):
    """3+ 次工具调用，f1 = 1.5"""
    metrics = ContextMetrics(tool_calls=3, output_tokens=100, turn_count=1)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.5)

    metrics = ContextMetrics(tool_calls=10, output_tokens=100, turn_count=1)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.5)


def test_low_output_tokens_no_boost(calc):
    """<=500 output tokens，f2 = 1.0"""
    metrics = ContextMetrics(tool_calls=0, output_tokens=500, turn_count=1)
    budget = calc.calculate(metrics)
    assert budget == ThinkingBudgetCalculator.BASE_BUDGET


def test_medium_output_tokens_small_boost(calc):
    """500 < output_tokens <= 2000，f2 = 1.1"""
    metrics = ContextMetrics(tool_calls=0, output_tokens=501, turn_count=1)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.1)

    metrics = ContextMetrics(tool_calls=0, output_tokens=2000, turn_count=1)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.1)


def test_high_output_tokens_large_boost(calc):
    """>2000 output tokens，f2 = 1.3"""
    metrics = ContextMetrics(tool_calls=0, output_tokens=2001, turn_count=1)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.3)


def test_low_turn_count_no_boost(calc):
    """<=5 轮对话，f3 = 1.0"""
    metrics = ContextMetrics(tool_calls=0, output_tokens=100, turn_count=5)
    budget = calc.calculate(metrics)
    assert budget == ThinkingBudgetCalculator.BASE_BUDGET


def test_medium_turn_count_small_boost(calc):
    """5 < turn_count <= 10，f3 = 1.1"""
    metrics = ContextMetrics(tool_calls=0, output_tokens=100, turn_count=6)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.1)

    metrics = ContextMetrics(tool_calls=0, output_tokens=100, turn_count=10)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.1)


def test_high_turn_count_large_boost(calc):
    """>10 轮对话，f3 = 1.2"""
    metrics = ContextMetrics(tool_calls=0, output_tokens=100, turn_count=11)
    budget = calc.calculate(metrics)
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.2)


# -------- 多因子组合测试 --------

def test_all_factors_max(calc):
    """三因子同时最大：1.5 × 1.3 × 1.2 = 2.34"""
    metrics = ContextMetrics(tool_calls=3, output_tokens=2001, turn_count=11)
    budget = calc.calculate(metrics)
    expected = int(ThinkingBudgetCalculator.BASE_BUDGET * 1.5 * 1.3 * 1.2)
    assert budget == expected
    # 应在 [MIN, MAX] 区间内
    assert ThinkingBudgetCalculator.MIN <= budget <= ThinkingBudgetCalculator.MAX


def test_all_factors_medium(calc):
    """三因子中等：1.2 × 1.1 × 1.1 = 1.452"""
    metrics = ContextMetrics(tool_calls=1, output_tokens=501, turn_count=6)
    budget = calc.calculate(metrics)
    expected = int(ThinkingBudgetCalculator.BASE_BUDGET * 1.2 * 1.1 * 1.1)
    assert budget == expected


# -------- Clamp 边界测试 --------

def test_clamp_to_min(calc):
    """当所有因子都为 1.0 且 BASE_BUDGET 在 [MIN, MAX] 内时，结果 = BASE_BUDGET"""
    # 这个 case 实际上不会被 clamp，因为 BASE=10000 在 [2000, 50000] 内
    metrics = ContextMetrics.initial()
    budget = calc.calculate(metrics)
    assert budget >= ThinkingBudgetCalculator.MIN


def test_clamp_to_max(calc):
    """极端高复杂度场景，结果不应超过 MAX。

    注：当前因子最大乘积 = 1.5 × 1.3 × 1.2 = 2.34，对应预算 23400，
    不会触及 MAX=50000。MAX 是安全上限，防止未来扩展因子时失控。
    本测试验证极端场景下结果不超过 MAX。
    """
    metrics = ContextMetrics(tool_calls=100, output_tokens=10000, turn_count=100)
    budget = calc.calculate(metrics)
    # 当前因子下最大值为 23400
    assert budget == int(ThinkingBudgetCalculator.BASE_BUDGET * 1.5 * 1.3 * 1.2)
    assert budget <= ThinkingBudgetCalculator.MAX


def test_min_value_is_2000():
    assert ThinkingBudgetCalculator.MIN == 2000


def test_max_value_is_50000():
    assert ThinkingBudgetCalculator.MAX == 50000


def test_base_value_is_10000():
    assert ThinkingBudgetCalculator.BASE_BUDGET == 10000


# -------- budget_to_effort 映射测试 --------

def test_budget_to_effort_low():
    """2000-10000 → low"""
    assert budget_to_effort(2000) == "low"
    assert budget_to_effort(5000) == "low"
    assert budget_to_effort(10000) == "medium"  # 边界：10000 属于 medium


def test_budget_to_effort_medium():
    """10000-25000 → medium"""
    assert budget_to_effort(10000) == "medium"
    assert budget_to_effort(15000) == "medium"
    assert budget_to_effort(25000) == "high"  # 边界：25000 属于 high


def test_budget_to_effort_high():
    """25000-50000 → high"""
    assert budget_to_effort(25000) == "high"
    assert budget_to_effort(40000) == "high"
    assert budget_to_effort(50000) == "max"  # 边界：50000 属于 max


def test_budget_to_effort_max():
    """>=50000 → max"""
    assert budget_to_effort(50000) == "max"
    assert budget_to_effort(60000) == "max"
    assert budget_to_effort(100000) == "max"


def test_budget_to_effort_below_min_clamps_to_low():
    """低于 MIN 的预算映射到 low"""
    assert budget_to_effort(0) == "low"
    assert budget_to_effort(1000) == "low"
    assert budget_to_effort(1999) == "low"


# -------- ContextMetrics dataclass 测试 --------

def test_context_metrics_initial():
    """ContextMetrics.initial() 应返回所有 factor=1.0 的默认值"""
    m = ContextMetrics.initial()
    assert m.tool_calls == 0
    assert m.output_tokens == 0
    assert m.turn_count == 1  # 首轮 turn_count=1


def test_context_metrics_to_dict():
    """to_dict 应正确序列化"""
    m = ContextMetrics(tool_calls=2, output_tokens=1000, turn_count=5)
    d = m.to_dict()
    assert d == {"tool_calls": 2, "output_tokens": 1000, "turn_count": 5}


def test_calculate_returns_int(calc):
    """calculate 应返回 int 类型"""
    metrics = ContextMetrics.initial()
    budget = calc.calculate(metrics)
    assert isinstance(budget, int)


def test_calculate_with_negative_values(calc):
    """负值应被视为 0"""
    metrics = ContextMetrics(tool_calls=-1, output_tokens=-100, turn_count=0)
    budget = calc.calculate(metrics)
    # 负值按 0 处理，所有 factor = 1.0
    assert budget == ThinkingBudgetCalculator.BASE_BUDGET
