"""
Dashscope模型调用封装模块
提供统一的模型调用接口，支持重试机制和错误处理
"""

import os
import time
from dotenv import load_dotenv
import dashscope
from http import HTTPStatus
from config import MODEL_CONFIG, get_api_key

# 加载环境变量并设置API Key
load_dotenv()
dashscope.api_key = get_api_key()


class DashscopeModel:
    """
    Dashscope模型调用封装类
    
    功能：
    1. 封装Dashscope SDK的调用细节
    2. 提供自动重试机制
    3. 统一的返回格式
    4. 错误处理和日志记录
    """
    
    def __init__(self, model_name=None, temperature=None):
        """
        初始化模型客户端
        
        Args:
            model_name: 模型名称，默认从配置读取
            temperature: 温度参数，默认从配置读取
        """
        self.model_name = model_name or MODEL_CONFIG["model_name"]
        self.temperature = temperature if temperature is not None else MODEL_CONFIG["temperature"]
        self.max_retries = MODEL_CONFIG["max_retries"]
        self.timeout = MODEL_CONFIG["timeout"]
        
    def call(self, prompt, retry_count=0):
        """
        调用模型并返回结果
        
        Args:
            prompt: 输入的提示词
            retry_count: 当前重试次数（内部使用）
            
        Returns:
            dict: 包含以下字段的字典
                - success: 是否成功
                - content: 模型输出的文本内容（成功时）
                - tokens: 消耗的token总数（成功时）
                - error: 错误信息（失败时）
        """
        try:
            # 调用Dashscope API
            response = dashscope.Generation.call(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                temperature=self.temperature,
                result_format="message"
            )
            
            # 检查响应状态
            if response.status_code == HTTPStatus.OK:
                # 提取输出内容
                content = response.output.choices[0].message.content
                tokens = response.usage.total_tokens if hasattr(response, 'usage') else 0
                
                return {
                    "success": True,
                    "content": content,
                    "tokens": tokens
                }
            else:
                # API返回错误，尝试重试
                error_msg = f"API错误: {response.code} - {response.message}"
                if retry_count < self.max_retries:
                    print(f"  ⚠️ 调用失败，{retry_count + 1}/{self.max_retries}次重试...")
                    time.sleep(1)  # 等待1秒后重试
                    return self.call(prompt, retry_count + 1)
                
                return {
                    "success": False,
                    "error": error_msg
                }
                
        except Exception as e:
            # 发生异常，尝试重试
            error_msg = f"调用异常: {str(e)}"
            if retry_count < self.max_retries:
                print(f"  ⚠️ 发生异常，{retry_count + 1}/{self.max_retries}次重试...")
                time.sleep(1)
                return self.call(prompt, retry_count + 1)
            
            return {
                "success": False,
                "error": error_msg
            }
    
    def call_with_retry(self, prompt):
        """
        带重试机制的模型调用（call的别名，保持向后兼容）
        
        Args:
            prompt: 输入的提示词
            
        Returns:
            dict: 调用结果
        """
        return self.call(prompt)


class MockModel:
    """
    模拟模型（用于测试，不消耗API额度）
    """
    
    def __init__(self, model_name="mock-model", temperature=0.1):
        self.model_name = model_name
        self.temperature = temperature
        
    def call(self, prompt):
        """模拟调用，返回预设的测试数据"""
        import random
        
        # 模拟不同的响应
        responses = [
            "这是一个模拟的回答。在实际运行中，这里会显示模型的真实输出。",
            "模拟思维链输出：\n1. 首先分析问题...\n2. 然后逐步推理...\n3. 最终得出结论...\n\n最终答案：示例答案",
            "根据提示词的要求，我进行了详细的思考...\n\n最终答案是：42"
        ]
        
        return {
            "success": True,
            "content": random.choice(responses),
            "tokens": random.randint(50, 200)
        }


# 便捷函数：快速创建模型实例
def create_model(use_mock=False, **kwargs):
    """
    创建模型实例的工厂函数
    
    Args:
        use_mock: 是否使用模拟模型（测试用）
        **kwargs: 传递给模型构造函数的参数
        
    Returns:
        DashscopeModel 或 MockModel 实例
    """
    if use_mock:
        return MockModel(**kwargs)
    return DashscopeModel(**kwargs)