#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
DeepSeek API Token使用情况分析工具
"""

import os
import sys
import argparse
import json
from typing import Dict, Any
import pandas as pd
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from api.deepseek_api import DeepSeekAPI
from utils.config_loader import ConfigLoader


def get_deepseek_api() -> DeepSeekAPI:
    """
    获取已配置的DeepSeekAPI实例
    """
    config_loader = ConfigLoader()
    deepseek_config = config_loader.get_deepseek_config()
    
    if not deepseek_config or 'api_key' not in deepseek_config:
        print("错误: 未找到DeepSeek API配置或API密钥")
        sys.exit(1)
    
    return DeepSeekAPI(
        api_key=deepseek_config['api_key'],
        model=deepseek_config.get('model', 'deepseek-moe-16b-chat')
    )


def show_token_usage_stats():
    """
    显示Token使用统计信息
    """
    api = get_deepseek_api()
    stats = api.get_token_usage_stats()
    
    print("\n===== DeepSeek API Token使用统计 =====")
    print(f"总输入Tokens: {stats['total_prompt_tokens']:,}")
    print(f"总输出Tokens: {stats['total_completion_tokens']:,}")
    print(f"总Tokens: {stats['total_tokens']:,}")
    print(f"API调用次数: {stats['total_calls']}")
    print(f"平均每次调用Tokens: {stats['average_tokens_per_call']:.2f}")
    
    if stats['total_calls'] > 0:
        print("\n最近调用记录:")
        for call in stats['detailed_calls']:
            print(f"  - {call['timestamp']} | {call['symbol']} | {call['task_type']} | "
                  f"输入: {call['prompt_tokens']} | 输出: {call['completion_tokens']} | "
                  f"总计: {call['total_tokens']} | 耗时: {call['duration_seconds']}秒")
    
    print("\n注意: 这些统计数据仅反映当前程序运行期间的使用情况")


def export_token_usage(format_type: str = "json", output_path: str = None):
    """
    导出Token使用数据
    
    Args:
        format_type: 导出格式，'json'或'csv'
        output_path: 输出文件路径
    """
    api = get_deepseek_api()
    
    if output_path is None:
        # 如果未指定路径，则使用默认路径
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = f"deepseek_token_usage_{timestamp}.{format_type}"
    
    file_path = api.export_token_usage(output_path, format_type)
    
    if file_path:
        print(f"\nToken使用数据已导出到: {file_path}")
    else:
        print("\n导出Token使用数据失败")


def analyze_token_usage(json_file: str = None):
    """
    分析Token使用数据并显示统计
    
    Args:
        json_file: JSON格式的Token使用数据文件
    """
    try:
        # 加载数据
        if json_file and os.path.exists(json_file):
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
        else:
            # 使用当前内存中的数据
            api = get_deepseek_api()
            data = api.token_usage
        
        if not data or not data.get('calls'):
            print("没有可用的Token使用数据进行分析")
            return
        
        # 转换为DataFrame进行分析
        df = pd.DataFrame(data['calls'])
        
        # 按任务类型分析
        task_analysis = df.groupby('task_type').agg({
            'prompt_tokens': ['sum', 'mean'],
            'completion_tokens': ['sum', 'mean'],
            'total_tokens': ['sum', 'mean', 'count']
        })
        
        # 按符号(交易对)分析
        symbol_analysis = df.groupby('symbol').agg({
            'total_tokens': ['sum', 'mean', 'count']
        })
        
        # 计算总体统计
        total_tokens = data['total_tokens']
        total_calls = len(data['calls'])
        
        # 显示分析结果
        print("\n===== DeepSeek API Token使用分析 =====")
        print(f"总调用次数: {total_calls}")
        print(f"总Token使用: {total_tokens:,}")
        print(f"平均每次调用Token: {total_tokens / total_calls if total_calls else 0:.2f}")
        
        print("\n--- 按任务类型分析 ---")
        for task, stats in task_analysis.iterrows():
            calls = stats[('total_tokens', 'count')]
            total = stats[('total_tokens', 'sum')]
            avg = stats[('total_tokens', 'mean')]
            print(f"{task}: {calls}次调用, 共{total:,}tokens, 平均{avg:.2f}tokens/次")
        
        print("\n--- 按交易对分析 ---")
        for symbol, stats in symbol_analysis.iterrows():
            calls = stats[('total_tokens', 'count')]
            total = stats[('total_tokens', 'sum')]
            avg = stats[('total_tokens', 'mean')]
            print(f"{symbol}: {calls}次调用, 共{total:,}tokens, 平均{avg:.2f}tokens/次")
        
    except Exception as e:
        print(f"分析Token使用数据时出错: {e}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="DeepSeek API Token使用情况分析工具")
    
    subparsers = parser.add_subparsers(dest="command", help="命令")
    
    # 显示统计信息
    show_parser = subparsers.add_parser("show", help="显示Token使用统计")
    
    # 导出使用数据
    export_parser = subparsers.add_parser("export", help="导出Token使用数据")
    export_parser.add_argument(
        "--format", "-f", 
        choices=["json", "csv"], 
        default="json",
        help="导出格式 (默认: json)"
    )
    export_parser.add_argument(
        "--output", "-o",
        help="输出文件路径"
    )
    
    # 分析使用数据
    analyze_parser = subparsers.add_parser("analyze", help="分析Token使用数据")
    analyze_parser.add_argument(
        "--input", "-i",
        help="输入JSON文件路径，如果未指定则使用当前内存中的数据"
    )
    
    args = parser.parse_args()
    
    if args.command == "show":
        show_token_usage_stats()
    elif args.command == "export":
        export_token_usage(args.format, args.output)
    elif args.command == "analyze":
        analyze_token_usage(args.input)
    else:
        # 默认显示统计信息
        show_token_usage_stats()


if __name__ == "__main__":
    main() 