import logging
from openai import OpenAI
import os
import sys
import time
import random
from typing import Dict, Any, Optional

# 设置日志记录
logger = logging.getLogger(__name__)

# 默认配置
DEFAULT_API_KEYS = [
    "28cfe71a-c6fa-4c5d-9b4e-d8474f0d3b93",  # 默认测试密钥
]

DEFAULT_MODELS = {
    "online_bot": "bot-20250325102825-h9kpq",      # 联网智能体ID
    "offline_model": "ep-20250326090920-v7wns",    # 免费推理50W
    "doubao": "doubao-1-5-pro-32k-250115",         # 豆包模型
    "deepseek": "deepseek-r1-250120",              # Deepseek模型
}

# 获取项目根目录的绝对路径
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# 尝试导入配置
try:
    # 添加项目根目录到 Python 路径
    sys.path.append(os.path.dirname(ROOT_DIR))
    sys.path.append(ROOT_DIR)
    
    # 导入配置
    try:
        from scripts.config import get_random_api_key, get_model
        logger.info("成功从scripts.config导入配置")
    except ImportError:
        try:
            from src.scripts.config import get_random_api_key, get_model
            logger.info("成功从src.scripts.config导入配置")
        except ImportError:
            logger.warning("无法导入配置模块，使用默认配置")
            
            # 使用默认配置的实现
            def get_random_api_key():
                return random.choice(DEFAULT_API_KEYS)
                
            def get_model(model_type):
                return DEFAULT_MODELS.get(model_type, DEFAULT_MODELS["offline_model"])
except Exception as e:
    logger.error(f"导入配置时出错: {str(e)}，使用默认配置")
    
    # 使用默认配置的实现
    def get_random_api_key():
        return random.choice(DEFAULT_API_KEYS)
        
    def get_model(model_type):
        return DEFAULT_MODELS.get(model_type, DEFAULT_MODELS["offline_model"])

class ChatBot:
    def __init__(self, model_type: str = "online_bot"):
        """初始化聊天机器人
        
        Args:
            model_type: 要使用的模型类型，默认为联网智能体
        """
        try:
            # 从配置获取API密钥
            self.api_key = get_random_api_key()
            # 从配置获取模型ID
            self.model = get_model(model_type)
            
            logger.info(f"初始化ChatBot，使用模型: {self.model}")
            
            # 初始化OpenAI客户端
            self.client = OpenAI(
                base_url="https://ark.cn-beijing.volces.com/api/v3/bots",
                api_key=self.api_key
            )
            
            # 系统提示语
            self.system_message = {
                "role": "system",
                "content": """你是一个专业的股票分析助手，擅长进行深入的基本面分析。你的分析应该：

                1. 专业严谨
                - 使用准确的专业术语
                - 引用可靠的数据来源
                - 分析逻辑清晰
                - 结论有理有据
                
                2. 全面细致
                - 深入分析问题的各个方面
                - 关注细节和关键信息
                - 考虑多个影响因素
                - 提供详实的论据支持
                
                3. 客观中立
                - 保持独立判断
                - 不夸大或贬低
                - 平衡利弊分析
                - 指出潜在风险
                
                4. 实用性强
                - 分析结论具体明确
                - 建议具有可操作性
                - 关注实际投资价值
                - 提供清晰的决策参考
                
                5. 及时更新
                - 关注最新信息
                - 反映市场变化
                - 动态调整分析
                - 保持信息时效性
                
                请根据用户的具体需求，提供专业、深入的分析。如果遇到不确定的信息，请明确说明。"""
            }
            
            # 对话历史
            self.conversation_history = [self.system_message]
        except Exception as e:
            logger.error(f"初始化ChatBot时出错: {str(e)}")
            raise
        
    def format_reference(self, ref):
        """格式化参考资料"""
        try:
            if isinstance(ref, str):
                return ref
            elif isinstance(ref, dict):
                title = ref.get('title', '')
                summary = ref.get('summary', '')
                url = ref.get('url', '')
                publish_time = ref.get('publish_time', '')
                formatted_ref = []
                if title:
                    formatted_ref.append(f"标题：{title}")
                if summary:
                    formatted_ref.append(f"摘要：{summary}")
                if url:
                    formatted_ref.append(f"链接：{url}")
                if publish_time:
                    formatted_ref.append(f"发布时间：{publish_time}")
                
                return "\n".join(formatted_ref)
            else:
                return str(ref)
        except Exception as e:
            logger.error(f"格式化参考资料时出错: {str(e)}")
            return str(ref)
        
    def chat(self, user_input: str, temperature: float = 0.7, top_p: float = 0.7, max_tokens: int = 4096, frequency_penalty: float = 0.0) -> Dict[str, Any]:
        """与AI进行对话
        
        Args:
            user_input: 用户输入的问题
            temperature: 控制输出的随机性，范围0-2，默认0.7
            top_p: 控制输出的多样性，范围0-1，默认0.7
            max_tokens: 控制输出的最大长度，默认4096
            frequency_penalty: 控制重复惩罚，范围-2到2，默认0.0
            
        Returns:
            Dict[str, Any]: 包含以下字段的字典：
                - response: AI的回答内容
                - reasoning_process: 推理过程
                - references: 参考资料列表
        """
        try:
            # 添加用户消息到对话历史
            self.conversation_history.append({
                "role": "user",
                "content": user_input
            })
            
            # 调用API（流式输出）
            stream = self.client.chat.completions.create(
                model=self.model,
                messages=self.conversation_history,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                frequency_penalty=frequency_penalty,
                stream=True
            )
            
            # 收集完整的回复
            full_response = ""
            print("\nAI: ", end="", flush=True)
            
            # 用于存储参考资料和推理过程
            references = []
            reasoning_parts = []
            
            for chunk in stream:
                # 检查是否有参考资料
                if hasattr(chunk, "references"):
                    refs = chunk.references
                    if refs:
                        references.extend(refs)
                
                if not chunk.choices:
                    continue
                
                # 输出内容
                if chunk.choices[0].delta.content:
                    content = chunk.choices[0].delta.content
                    print(content, end="", flush=True)
                    full_response += content
                    time.sleep(0.02)
                
                # 收集推理过程
                elif hasattr(chunk.choices[0].delta, "reasoning_content"):
                    reasoning = chunk.choices[0].delta.reasoning_content
                    if reasoning:
                        reasoning_parts.append(reasoning)
            
            # 如果有推理过程，格式化并显示
            if reasoning_parts:
                print("\n\n推理过程：")
                reasoning_text = "".join(reasoning_parts)
                print(reasoning_text)
                print("-" * 50)

            # 如果有参考资料，格式化并显示
            if references:
                print("\n参考资料：")
                for ref in references:
                    formatted_ref = self.format_reference(ref)
                    print(f"\n{formatted_ref}")
                    print("-" * 50)
            
            # 从回答中提取推理过程和参考资料
            main_response = full_response
            reasoning_process = ""
            
            if "推理过程：" in full_response:
                parts = full_response.split("推理过程：")
                main_response = parts[0].strip()
                if len(parts) > 1:
                    reasoning_parts = parts[1].split("参考资料：")
                    reasoning_process = reasoning_parts[0].strip()
            
            # 添加AI回复到对话历史
            self.conversation_history.append({
                "role": "assistant",
                "content": main_response,
                "reasoning_process": reasoning_text,
                "references": references
            })
            
            return {
                "response": main_response,
                "reasoning_process": reasoning_text,
                "references": references
            }
            
        except Exception as e:
            logger.error(f"对话失败: {str(e)}")
            return {
                "response": f"抱歉，处理您的请求时出现错误: {str(e)}",
                "reasoning_process": "",
                "references": []
            }
    
    def clear_history(self):
        """清除对话历史"""
        self.conversation_history = [self.system_message]
        print("对话历史已清除")
    
    def run(self):
        """运行聊天机器人"""
        print("欢迎使用AI助手！输入 'quit' 退出，输入 'clear' 清除对话历史。")
        print("该版本支持联网搜索，可以回答实时信息。")
        print("-" * 50)
        
        while True:
            try:
                # 获取用户输入
                user_input = input("\n你: ").strip()
                
                # 检查退出命令
                if user_input.lower() == 'quit':
                    print("感谢使用，再见！")
                    break
                
                # 检查清除历史命令
                if user_input.lower() == 'clear':
                    self.clear_history()
                    continue
                
                # 如果输入为空，继续下一轮
                if not user_input:
                    continue
                
                # 获取AI回复
                self.chat(user_input)
                print("-" * 50)
                
            except KeyboardInterrupt:
                print("\n感谢使用，再见！")
                break
            except Exception as e:
                logger.error(f"运行错误: {str(e)}")
                print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    # 设置日志级别
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建并运行聊天机器人
    bot = ChatBot()
    bot.run() 