import os
import json
import requests
from typing import Dict, Any, List, Optional, Tuple
import time
import logging
import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("deepseek_token_usage.log"),
        logging.StreamHandler()
    ]
)

class DeepSeekAPI:
    """DeepSeek API交互类，用于进行大语言模型调用"""
    
    def __init__(self, api_key: str, model: str = "deepseek-moe-16b-chat"):
        """
        初始化DeepSeek API
        
        Args:
            api_key: DeepSeek API密钥
            model: 使用的模型名称
        """
        self.api_key = api_key
        self.model = model
        self.base_url = "https://api.deepseek.com/v1"
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        
        # Token 使用统计
        self.token_usage = {
            "total_prompt_tokens": 0,
            "total_completion_tokens": 0,
            "total_tokens": 0,
            "calls": []
        }
        
        # 创建日志记录器
        self.logger = logging.getLogger("DeepSeekAPI")
    
    def call_model(self, prompt: str, system_prompt: str = None, task_type: str = "未知任务", symbol: str = "未知", temperature: float = 0.2, max_tokens: int = 2000) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """
        调用DeepSeek大语言模型
        
        Args:
            prompt: 用户提示词
            system_prompt: 系统提示词，如果为None则使用默认值
            task_type: 任务类型，用于记录
            symbol: 交易对符号，用于记录
            temperature: 采样温度，控制输出随机性
            max_tokens: 最大生成token数
            
        Returns:
            (API响应, token使用信息)
        """
        if system_prompt is None:
            system_prompt = "你是一个专业的加密货币分析助手，擅长分析市场趋势、预测价格走向和提供交易建议。请始终使用中文回复，并确保输出格式规范的JSON。"
            
        usage_info = {}
        
        try:
            endpoint = f"{self.base_url}/chat/completions"
            
            payload = {
                "model": self.model,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}
                ],
                "temperature": temperature,
                "max_tokens": max_tokens
            }
            
            start_time = time.time()
            response = requests.post(endpoint, headers=self.headers, json=payload)
            response.raise_for_status()
            response_data = response.json()
            end_time = time.time()
            
            # 记录token使用情况
            if 'usage' in response_data:
                prompt_tokens = response_data['usage'].get('prompt_tokens', 0)
                completion_tokens = response_data['usage'].get('completion_tokens', 0)
                total_tokens = response_data['usage'].get('total_tokens', 0)
                
                usage_info = {
                    "prompt_tokens": prompt_tokens,
                    "completion_tokens": completion_tokens,
                    "total_tokens": total_tokens,
                    "task_type": task_type,
                    "symbol": symbol,
                    "model": self.model,
                    "timestamp": datetime.datetime.now().isoformat(),
                    "duration_seconds": round(end_time - start_time, 2)
                }
                
                # 更新总计
                self.token_usage["total_prompt_tokens"] += prompt_tokens
                self.token_usage["total_completion_tokens"] += completion_tokens
                self.token_usage["total_tokens"] += total_tokens
                self.token_usage["calls"].append(usage_info)
                
                # 记录到日志
                self.logger.info(
                    f"DeepSeek API调用 - 任务: {task_type}, 符号: {symbol}, "
                    f"输入tokens: {prompt_tokens}, 输出tokens: {completion_tokens}, "
                    f"总tokens: {total_tokens}, 耗时: {round(end_time - start_time, 2)}秒"
                )
            
            return response_data, usage_info
        
        except Exception as e:
            error_msg = f"调用DeepSeek API时出错: {e}"
            self.logger.error(error_msg)
            return {}, usage_info
    
    def extract_json_from_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
        """
        从响应中提取JSON数据
        
        Args:
            response: API响应
            
        Returns:
            提取的JSON数据
        """
        try:
            if 'choices' in response and len(response['choices']) > 0:
                content = response['choices'][0]['message']['content']
                
                # 尝试从响应中提取JSON
                start_idx = content.find('{')
                end_idx = content.rfind('}') + 1
                
                if start_idx != -1 and end_idx != -1:
                    json_str = content[start_idx:end_idx]
                    return json.loads(json_str)
                
                return {"error": "无法从响应中提取JSON", "raw_content": content}
            
            return {"error": "API响应格式不正确", "raw_response": response}
        
        except Exception as e:
            error_msg = f"解析响应时出错: {e}"
            self.logger.error(error_msg)
            return {"error": str(e), "raw_response": response}
    
    def get_token_usage_stats(self) -> Dict[str, Any]:
        """
        获取Token使用统计信息
        
        Returns:
            包含使用统计的字典
        """
        return {
            "total_prompt_tokens": self.token_usage["total_prompt_tokens"],
            "total_completion_tokens": self.token_usage["total_completion_tokens"],
            "total_tokens": self.token_usage["total_tokens"],
            "total_calls": len(self.token_usage["calls"]),
            "average_tokens_per_call": self.token_usage["total_tokens"] / len(self.token_usage["calls"]) if self.token_usage["calls"] else 0,
            "detailed_calls": self.token_usage["calls"][-10:]  # 仅返回最近10次调用详情
        }
    
    def export_token_usage(self, file_path: str = None, format: str = "json") -> str:
        """
        导出Token使用数据到文件
        
        Args:
            file_path: 文件路径，如果为None则自动生成
            format: 导出格式，支持'json'或'csv'
            
        Returns:
            导出文件的路径
        """
        if file_path is None:
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            file_path = f"deepseek_token_usage_{timestamp}.{format}"
        
        try:
            if format.lower() == "json":
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(self.token_usage, f, indent=2, ensure_ascii=False)
            elif format.lower() == "csv":
                import csv
                
                with open(file_path, 'w', newline='', encoding='utf-8') as f:
                    writer = csv.writer(f)
                    # 写入表头
                    writer.writerow([
                        "timestamp", "task_type", "symbol", "model", 
                        "prompt_tokens", "completion_tokens", "total_tokens", 
                        "duration_seconds"
                    ])
                    
                    # 写入数据
                    for call in self.token_usage["calls"]:
                        writer.writerow([
                            call.get("timestamp", ""),
                            call.get("task_type", ""),
                            call.get("symbol", ""),
                            call.get("model", ""),
                            call.get("prompt_tokens", 0),
                            call.get("completion_tokens", 0),
                            call.get("total_tokens", 0),
                            call.get("duration_seconds", 0)
                        ])
                    
                    # 写入总计
                    writer.writerow([])
                    writer.writerow([
                        f"总计 (调用次数: {len(self.token_usage['calls'])})",
                        "", "", "",
                        self.token_usage["total_prompt_tokens"],
                        self.token_usage["total_completion_tokens"],
                        self.token_usage["total_tokens"],
                        ""
                    ])
            else:
                raise ValueError(f"不支持的格式: {format}，仅支持 'json' 或 'csv'")
                
            self.logger.info(f"Token使用数据已导出到: {file_path}")
            return file_path
            
        except Exception as e:
            error_msg = f"导出Token使用数据时出错: {e}"
            self.logger.error(error_msg)
            return "" 