#!/usr/bin/env python3
"""
Token and Cost Analysis Tool - 分析使用全部工具的成本影响
"""

import os
import sys
sys.path.append(os.getcwd())

import json
from slack_bot.llm.gemini import get_system_instruction
from slack_bot.tools.groups import get_tool_preset

print("=" * 80)
print("Token 和成本分析工具")
print("=" * 80)

# 估算 token 数量（粗略估算：1 token ≈ 3 chars）
def estimate_tokens(text):
    return len(text) // 3

# 获取数据
system_prompt = get_system_instruction()

print(f"\n📋 基础配置:")
print(f"  System Prompt: {len(system_prompt):,} chars (~{estimate_tokens(system_prompt):,} tokens)")

# 分析不同模式的工具定义大小
print(f"\n🔧 工具定义大小:")

modes = ["light", "standard", "full", "all"]
for mode in modes:
    tools, funcs = get_tool_preset(mode)
    tools_json = json.dumps(tools, ensure_ascii=False)

    print(f"\n  {mode.upper()} 模式:")
    print(f"    工具数量: {len(tools)}")
    print(f"    Schema 大小: {len(tools_json):,} chars (~{estimate_tokens(tools_json):,} tokens)")

    # 列出工具名称
    tool_names = [t['function']['name'] for t in tools]
    print(f"    工具列表: {', '.join(tool_names)}")

# 模拟典型请求的 token 消耗
print(f"\n💬 典型请求 Token 消耗估算:")

test_scenarios = [
    {
        "name": "简单查询（无上下文）",
        "user_message": "今天睡得怎么样？",
        "context_msgs": 0,
        "avg_context_chars": 0
    },
    {
        "name": "正常对话（10条历史）",
        "user_message": "昨天的步数",
        "context_msgs": 10,
        "avg_context_chars": 500  # 每条消息平均 500 字符
    },
    {
        "name": "长对话（30条历史）",
        "user_message": "过去一个月的HRV趋势",
        "context_msgs": 30,
        "avg_context_chars": 500
    },
    {
        "name": "复杂查询（带工具结果）",
        "user_message": "分析一下为什么最近睡眠不好",
        "context_msgs": 5,
        "avg_context_chars": 500,
        "tool_result_chars": 3000  # 工具返回的数据
    }
]

for mode in ["light", "all"]:
    tools, _ = get_tool_preset(mode)
    tools_json = json.dumps(tools, ensure_ascii=False)
    tools_tokens = estimate_tokens(tools_json)

    print(f"\n  === {mode.upper()} 模式 ===")

    for scenario in test_scenarios:
        # 计算输入 tokens
        input_tokens = (
            estimate_tokens(system_prompt) +  # System prompt
            tools_tokens +  # Tools schema
            estimate_tokens(scenario["user_message"]) +  # User message
            (scenario["context_msgs"] * estimate_tokens("X" * scenario["avg_context_chars"]))  # Context
        )

        if "tool_result_chars" in scenario:
            input_tokens += estimate_tokens("X" * scenario["tool_result_chars"])

        # 估算输出 tokens（假设平均响应 500 字符）
        output_tokens = estimate_tokens("X" * 500)

        total_tokens = input_tokens + output_tokens

        print(f"\n    {scenario['name']}:")
        print(f"      输入: ~{input_tokens:,} tokens")
        print(f"      输出: ~{output_tokens:,} tokens")
        print(f"      总计: ~{total_tokens:,} tokens")

# 成本计算
print(f"\n{'=' * 80}")
print("💰 成本估算")
print(f"{'=' * 80}")

# OpenRouter Gemini 价格（2024年价格）
INPUT_PRICE = 0.075 / 1_000_000  # $0.075 per 1M input tokens
OUTPUT_PRICE = 0.30 / 1_000_000  # $0.30 per 1M output tokens

print(f"\nOpenRouter Gemini 3 Flash 价格:")
print(f"  输入: $0.075 / 1M tokens")
print(f"  输出: $0.30 / 1M tokens")

# 场景分析
print(f"\n场景分析:")

daily_scenarios = [
    {
        "name": "轻度使用",
        "queries_per_day": 30,
        "avg_input_tokens": 4000,
        "avg_output_tokens": 200
    },
    {
        "name": "中度使用",
        "queries_per_day": 100,
        "avg_input_tokens": 5000,
        "avg_output_tokens": 300
    },
    {
        "name": "重度使用",
        "queries_per_day": 300,
        "avg_input_tokens": 6000,
        "avg_output_tokens": 500
    }
]

for scenario in daily_scenarios:
    daily_input = scenario["queries_per_day"] * scenario["avg_input_tokens"]
    daily_output = scenario["queries_per_day"] * scenario["avg_output_tokens"]

    monthly_input = daily_input * 30
    monthly_output = daily_output * 30

    monthly_cost = (monthly_input * INPUT_PRICE) + (monthly_output * OUTPUT_PRICE)
    monthly_cost_cny = monthly_cost * 7.2  # USD to CNY

    print(f"\n  {scenario['name']} ({scenario['queries_per_day']}次/天):")
    print(f"    月输入 tokens: {monthly_input:,}")
    print(f"    月输出 tokens: {monthly_output:,}")
    print(f"    月成本: ${monthly_cost:.2f} (~¥{monthly_cost_cny:.2f})")

# 对比 light vs all
print(f"\n{'=' * 80}")
print("📊 Light vs All 模式对比")
print(f"{'=' * 80}")

light_tools, _ = get_tool_preset("light")
all_tools, _ = get_tool_preset("all")

light_tools_json = json.dumps(light_tools, ensure_ascii=False)
all_tools_json = json.dumps(all_tools, ensure_ascii=False)

light_tools_tokens = estimate_tokens(light_tools_json)
all_tools_tokens = estimate_tokens(all_tools_json)

token_diff = all_tools_tokens - light_tools_tokens
token_diff_pct = (token_diff / light_tools_tokens) * 100

print(f"\nTools Schema Tokens:")
print(f"  Light (6 tools): {light_tools_tokens:,} tokens")
print(f"  All (21 tools): {all_tools_tokens:,} tokens")
print(f"  差异: +{token_diff:,} tokens (+{token_diff_pct:.1f}%)")

# 计算额外成本
print(f"\n额外成本计算（假设 100 次/天）:")

daily_queries = 100
monthly_queries = daily_queries * 30

# 每次请求的额外 token（只影响输入）
extra_cost_per_query = token_diff * INPUT_PRICE
monthly_extra_cost = extra_cost_per_query * monthly_queries
monthly_extra_cost_cny = monthly_extra_cost * 7.2

print(f"  每次查询额外成本: ${extra_cost_per_query:.6f}")
print(f"  月额外成本: ${monthly_extra_cost:.4f} (~¥{monthly_extra_cost_cny:.2f})")

# 结论
print(f"\n{'=' * 80}")
print("🎯 结论")
print(f"{'=' * 80}")

if monthly_extra_cost_cny < 5:
    print(f"\n✅ 使用 All 模式的额外成本很低 (¥{monthly_extra_cost_cny:.2f}/月)")
    print(f"   相比获得的完整功能，这个成本可以忽略不计")
    print(f"\n   推荐: 使用 All 模式，获得全部 21 个工具的能力")
elif monthly_extra_cost_cny < 20:
    print(f"\n⚠️  使用 All 模式的额外成本较低 (¥{monthly_extra_cost_cny:.2f}/月)")
    print(f"   如果需要完整功能，这个成本是可以接受的")
    print(f"\n   推荐: 根据实际需求选择")
else:
    print(f"\n⚠️  使用 All 模式的额外成本较高 (¥{monthly_extra_cost_cny:.2f}/月)")
    print(f"   建议评估是否真的需要所有工具")

print(f"\n{'=' * 80}")