"""
测试 Token 使用统计的索引功能
"""
import asyncio
import sys
from pathlib import Path

# 添加 backend 到路径
sys.path.insert(0, str(Path(__file__).parent / "backend"))

from backend.services.token_usage_service import token_usage_service
from backend.models.internal import TokenUsageStatus


async def test_token_usage_indexes():
    """测试 Token 使用统计的索引功能"""
    
    print("=" * 80)
    print("Token 使用统计 - 索引功能测试")
    print("=" * 80)
    
    # 1. 记录一些测试数据
    print("\n📝 记录测试数据...")
    
    test_records = [
        {
            "chat_id": "测试角色/测试聊天1",
            "role_name": "测试角色",
            "chat_name": "测试聊天1",
            "prompt_tokens": 100,
            "completion_tokens": 200,
            "total_tokens": 300,
            "api_url": "https://api.openai.com/v1"
        },
        {
            "chat_id": "测试角色/测试聊天2",
            "role_name": "测试角色",
            "chat_name": "测试聊天2",
            "prompt_tokens": 150,
            "completion_tokens": 250,
            "total_tokens": 400,
            "api_url": "https://api.deepseek.com/v1"
        },
        {
            "chat_id": "测试角色/测试聊天1",
            "role_name": "测试角色",
            "chat_name": "测试聊天1",
            "prompt_tokens": 120,
            "completion_tokens": 180,
            "total_tokens": 300,
            "api_url": "https://api.openai.com/v1"
        }
    ]
    
    for record_data in test_records:
        await token_usage_service.record_usage(**record_data)
        print(f"  ✅ 记录: {record_data['api_url']} - {record_data['total_tokens']} tokens")
    
    # 2. 测试 API URL 索引查询
    print("\n🔍 测试 API URL 索引查询...")
    api_url_stats = await token_usage_service.get_api_url_stats()
    
    if api_url_stats:
        print(f"\n  找到 {len(api_url_stats)} 个 API URL:")
        for url, stats in api_url_stats.items():
            print(f"\n  📊 {url}:")
            print(f"     - 总调用次数: {stats['count']}")
            print(f"     - 总 Token 数: {stats['totalTokens']}")
            print(f"     - Prompt Tokens: {stats['totalPromptTokens']}")
            print(f"     - Completion Tokens: {stats['totalCompletionTokens']}")
    else:
        print("  ⚠️  没有找到 API URL 统计数据")
    
    # 3. 测试每日索引查询
    print("\n🔍 测试每日索引查询...")
    from datetime import datetime
    now = datetime.now()
    daily_stats = await token_usage_service.get_daily_stats(now.year, now.month)
    
    if daily_stats:
        print(f"\n  找到 {len(daily_stats)} 天的数据:")
        for day, stats in list(daily_stats.items())[:5]:  # 只显示前5天
            print(f"  📅 {day}:")
            print(f"     - 调用次数: {stats['count']}")
            print(f"     - 总 Token 数: {stats['totalTokens']}")
    else:
        print("  ⚠️  没有找到每日统计数据")
    
    print("\n" + "=" * 80)
    print("✅ 测试完成!")
    print("=" * 80)


if __name__ == "__main__":
    asyncio.run(test_token_usage_indexes())