import os
from datetime import datetime

def save_usage_info_to_txt(usage_info, total_duration, money, base_dir, video_dir, save_dir="/root/autodl-tmp/video_llm", run_timestamp=None):
    """
    保存API使用情况信息到TXT文件
    :param usage_info: API使用情况对象
    :param total_duration: API总响应时间
    :param money: 费用信息字典
    :param video_path: 原视频文件路径
    :param video_dir: 视频目录名
    :param save_dir: 保存目录
    :param run_timestamp: 运行时间戳，用于创建统一的文件夹
    :return: 保存的文件路径
    """
    # 创建保存目录 - 每次运行创建新的时间戳文件夹
    video_name = os.path.splitext(os.path.basename(base_dir))[0]
    if run_timestamp is None:
        run_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    txt_dir = os.path.join(save_dir, "cost", video_name, run_timestamp)
    os.makedirs(txt_dir, exist_ok=True)
    
    # 生成文件名（只使用片段名，不包含时间戳）
    txt_filename = f"tokens_{video_dir}.txt"
    txt_path = os.path.join(txt_dir, txt_filename)
    
    # 格式化使用情况信息
    usage_content = f"""API使用情况统计
=====================================
视频文件: {video_dir}
统计时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
=====================================

📊 使用情况详情:
{usage_info}

⏱️ 时间统计:
- API总响应时间: {total_duration:.2f} 秒

💰 费用统计:
- 输出费用: ${money.get('output_momey', 0):.6f}
- 提示词费用: ${money.get('prompt_momey', 0):.6f}
- 视频费用: ${money.get('video_momey', 0):.6f}
- 音频费用: ${money.get('audio_momey', 0):.6f}
- 总费用: ${money.get('sum_momey', 0):.6f}

📈 Token统计:
- 总输入Token: {usage_info.total_tokens if hasattr(usage_info, 'total_tokens') else 'N/A'}
- 提示词Token: {usage_info.prompt_tokens if hasattr(usage_info, 'prompt_tokens') else 'N/A'}
- 输出Token: {usage_info.completion_tokens if hasattr(usage_info, 'completion_tokens') else 'N/A'}

🔍 详细Token信息:
- 文本Token: {usage_info.prompt_tokens_details.text_tokens if hasattr(usage_info.prompt_tokens_details, 'text_tokens') else 'N/A'}
- 视频Token: {usage_info.prompt_tokens_details.video_tokens if hasattr(usage_info.prompt_tokens_details, 'video_tokens') else 'N/A'}
- 音频Token: {usage_info.prompt_tokens_details.audio_tokens if hasattr(usage_info.prompt_tokens_details, 'audio_tokens') else 'N/A'}
"""
    
    # 保存到文件
    try:
        with open(txt_path, 'w', encoding='utf-8') as f:
            f.write(usage_content)
        print(f"\n✅ 使用情况已保存到: {txt_path}")
        return txt_path
    except Exception as e:
        print(f"\n❌ 保存使用情况文件失败: {e}")
        return None

def save_simple_usage_info(usage_info, total_duration, money, video_path, save_dir="/root/autodl-tmp/final_output"):
    """
    保存简化的API使用情况信息到TXT文件
    :param usage_info: API使用情况对象
    :param total_duration: API总响应时间
    :param money: 费用信息字典
    :param video_path: 原视频文件路径
    :param save_dir: 保存目录
    :return: 保存的文件路径
    """
    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)
    
    # 生成文件名
    video_name = os.path.splitext(os.path.basename(video_path))[0]
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    txt_filename = f"usage_{video_name}_{timestamp}.txt"
    txt_path = os.path.join(save_dir, txt_filename)
    
    # 简化的使用情况信息
    usage_content = f"""API使用情况统计
=====================================
视频文件: {video_path}
统计时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
=====================================

使用情况: {usage_info}
总耗时: {total_duration:.2f}秒
费用: {money}
"""
    
    # 保存到文件
    try:
        with open(txt_path, 'w', encoding='utf-8') as f:
            f.write(usage_content)
        print(f"\n✅ 使用情况已保存到: {txt_path}")
        return txt_path
    except Exception as e:
        print(f"\n❌ 保存使用情况文件失败: {e}")
        return None

if __name__ == "__main__":
    # 测试函数
    test_usage_info = "测试使用情况"
    test_duration = 10.5
    test_money = {
        'output_momey': 0.001,
        'prompt_momey': 0.002,
        'video_momey': 0.003,
        'audio_momey': 0.004,
        'sum_momey': 0.010
    }
    test_video_path = "/root/autodl-tmp/video/test.mp4"
    
    # 测试详细版本
    save_usage_info_to_txt(test_usage_info, test_duration, test_money, test_video_path, "test_segment")
    
    # 测试简化版本
    save_simple_usage_info(test_usage_info, test_duration, test_money, test_video_path)