import logging
from openai import OpenAI
import os
import sys
import time


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

# 获取项目根目录的绝对路径
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_model_config
        logger.info("成功从scripts.config导入配置")
    except ImportError:
        try:
            from src.scripts.config import get_model_config
            logger.info("成功从src.scripts.config导入配置")
        except ImportError:
            logger.warning("无法导入配置模块，使用默认配置")
            
            # 使用默认配置的实现
            def get_model_config(platform: str, model_type: str):
                return {
                    "base_url": "https://ark.cn-beijing.volces.com/api/v3/",
                    "api_key": "28cfe71a-c6fa-4c5d-9b4e-d8474f0d3b93",
                    "model": "ep-20250326090920-v7wns"
                }
except Exception as e:
    logger.error(f"导入配置时出错: {str(e)}，使用默认配置")
    
    # 使用默认配置的实现
    def get_model_config(platform: str, model_type: str):
        return {
            "base_url": "https://ark.cn-beijing.volces.com/api/v3/",
            "api_key": "28cfe71a-c6fa-4c5d-9b4e-d8474f0d3b93",
            "model": "ep-20250326090920-v7wns"
        }

class ChatBot:
    def __init__(self, platform: str = "volc", model_type: str = "offline_model"):
        """初始化聊天机器人
        
        Args:
            platform: 平台名称（volc/doubao/deepseek），默认为volc
            model_type: 要使用的模型类型，默认为offline_model
        """
        try:
            # 从配置获取API配置
            config = get_model_config(platform, model_type)
            self.api_key = config["api_key"]
            self.model = config["model"]
            
            logger.info(f"初始化ChatBot，使用平台: {platform}, 模型: {self.model}")
            
            # 初始化OpenAI客户端
            self.client = OpenAI(
                base_url=config["base_url"],
                api_key=self.api_key,
                timeout=600
            )
            
            # 系统提示语
            self.system_message = {
                "role": "system",
                "content": """你是一位经验丰富的专业投资经理，擅长基本面分析和投资决策。你的分析特点如下：
                            
                            1. 分析风格：
                            - 专业、客观、理性
                            - 注重数据支撑
                            - 关注风险控制
                            - 重视投资性价比
                            
                            2. 分析框架：
                            - 公司基本面分析
                            - 行业竞争格局
                            - 估值水平评估
                            - 风险因素识别
                            - 投资机会判断
                            
                            3. 输出要求：
                            - 简明扼要，重点突出
                            - 逻辑清晰，层次分明
                            - 数据准确，论据充分
                            - 结论明确，建议具体
                            
                            请用专业投资经理的视角，对股票进行深入分析和投资建议。如果信息不足，请明确指出。"""
            }
            
            # 对话历史
            self.conversation_history = [self.system_message]
        except Exception as e:
            logger.error(f"初始化ChatBot时出错: {str(e)}")
            raise

    def chat(self, user_input: str, temperature: float = 1.0, top_p: float = 0.7, max_tokens: int = 4096, frequency_penalty: float = 0.0) -> str:
        """处理用户输入并返回AI回复
        
        Args:
            user_input: 用户输入的问题
            temperature: 控制输出的随机性，范围0-2，默认1.0
            top_p: 控制输出的多样性，范围0-1，默认0.7
            max_tokens: 控制输出的最大长度，默认4096
            frequency_penalty: 控制重复惩罚，范围-2到2，默认0.0
            
        Returns:
            str: AI的回答内容
        """
        try:
            # # 添加用户消息到对话历史-多轮
            self.conversation_history.append({
                "role": "user",
                "content": user_input
            })
            extra_body = {"chat_template_kwargs": {"enable_thinking": True}}
            
            # 调用OpenAI 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,
                extra_body=extra_body,
                timeout=600
            )
            
            # 收集完整的回复
            full_response = ""
            print("\nAI: ", end="", flush=True)
            
            for chunk in stream:
                if not chunk.choices:
                    continue
                content = chunk.choices[0].delta.content
                if content:
                    print(content, end="", flush=True)
                    full_response += content
                    # 添加一个小延迟，使输出更自然
                    time.sleep(0.02)
            
            print()  # 换行
            
            # 添加AI回复到对话历史
            self.conversation_history.append({
                "role": "assistant",
                "content": full_response
            })
            
            return full_response
            
        except Exception as e:
            logger.error(f"聊天失败: {str(e)}")
            error_msg = f"抱歉，发生错误: {str(e)}"
            print(f"\n{error_msg}")
            return error_msg
    
    def clear_history(self):
        """清除对话历史"""
        self.conversation_history = [self.system_message]
        print("对话历史已清除")
    
    def run(self):
        """运行聊天机器人"""
        print("欢迎使用AI助手！输入 'quit' 退出，输入 'clear' 清除对话历史。")
        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() 