import logging
import os
import sys
import time
from typing import Dict, Any
from datetime import datetime

from openai import OpenAI

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

# 获取项目根目录
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# 导入配置
sys.path.append(os.path.dirname(ROOT_DIR))
sys.path.append(ROOT_DIR)

try:
    from scripts.config import get_random_api_key, get_model
except ImportError:
    from src.scripts.config import get_random_api_key, get_model


class ChatBot:
    def __init__(self, model_type: str = "online_bot", enable_web_search: bool = True):
        """初始化聊天机器人
        
        Args:
            model_type: 模型类型，默认为 online_bot
            enable_web_search: 是否启用联网搜索，默认True
        """
        self.api_key = get_random_api_key()
        self.model = get_model(model_type)
        self.enable_web_search = enable_web_search
        
        logger.info(f"初始化ChatBot，模型: {self.model}, 联网搜索: {self.enable_web_search}")
        
        # 初始化客户端
        self.client = OpenAI(
            base_url="https://ark.cn-beijing.volces.com/api/v3",
            api_key=self.api_key
        )
        
        # 联网搜索工具配置
        self.tools = [{
            "type": "web_search",
            "limit": 10,
            "sources": ["douyin", "moji", "toutiao"],
            "user_location": {
                "type": "approximate",
                "country": "中国",
                "region": "北京",
                "city": "北京"
            }
        }] if enable_web_search else None
        
        # 系统提示语
        self.system_message = """你是一个专业的股票分析助手，擅长进行深入的基本面分析。你的分析应该：

1. 专业严谨 - 使用准确的专业术语，引用可靠的数据来源，分析逻辑清晰，结论有理有据
2. 全面细致 - 深入分析问题的各个方面，关注细节和关键信息，考虑多个影响因素
3. 客观中立 - 保持独立判断，不夸大或贬低，平衡利弊分析，指出潜在风险
4. 实用性强 - 分析结论具体明确，建议具有可操作性，关注实际投资价值
5. 及时更新 - 关注最新信息，反映市场变化，动态调整分析，保持信息时效性

请根据用户的具体需求，提供专业、深入的分析。如果遇到不确定的信息，请明确说明。"""
        
        # 对话历史（简化版，只保存用户和助手消息）
        self.conversation_history = []
    
    def format_reference(self, ref):
        """格式化参考资料"""
        if isinstance(ref, str):
            return ref
        elif isinstance(ref, dict):
            parts = []
            if ref.get('title'):
                parts.append(f"标题：{ref['title']}")
            if ref.get('summary'):
                parts.append(f"摘要：{ref['summary']}")
            if ref.get('url'):
                parts.append(f"链接：{ref['url']}")
            if ref.get('publish_time'):
                parts.append(f"发布时间：{ref['publish_time']}")
            return "\n".join(parts)
        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, reasoning_process, references, tool_usage, tool_usage_details
        """
        try:
            # 添加用户消息到对话历史
            self.conversation_history.append({
                "role": "user",
                "content": [{"type": "input_text", "text": user_input}]
            })
            
            # 构建输入消息
            input_messages = [
                # 系统提示词
                {
                    "role": "system",
                    "content": [{"type": "input_text", "text": self.system_message}]
                }
            ]
            
            # 添加最近的对话历史（保留最近5轮）
            for msg in self.conversation_history[-5:]:
                input_messages.append(msg)
            
            # 调用API（流式输出）
            response = self.client.responses.create(
                model=self.model,
                input=input_messages,
                tools=self.tools,
                stream=True
            )
            
            # 收集回复
            full_response = ""
            reasoning_text = ""
            references = []
            tool_usage = None
            tool_usage_details = None
            search_keywords = []
            
            # 状态变量
            thinking_started = False
            answering_started = False
            
            print("\nAI: ", end="", flush=True)
            
            # 处理事件流
            for chunk in response:
                chunk_type = getattr(chunk, "type", "")
                
                # 处理AI思考过程
                if chunk_type == "response.reasoning_summary_text.delta":
                    if not thinking_started:
                        thinking_started = True
                    delta = getattr(chunk, "delta", "")
                    if delta:
                        reasoning_text += delta
                
                # 处理搜索状态
                elif "web_search_call" in chunk_type:
                    if "in_progress" in chunk_type:
                        logger.debug("开始搜索")
                    elif "completed" in chunk_type:
                        logger.debug("搜索完成")
                
                # 处理搜索关键词
                elif (chunk_type == "response.output_item.done" 
                      and hasattr(chunk, "item") 
                      and str(getattr(chunk.item, "id", "")).startswith("ws_")):
                    if hasattr(chunk.item, "action") and hasattr(chunk.item.action, "query"):
                        search_keyword = chunk.item.action.query
                        search_keywords.append(search_keyword)
                        logger.debug(f"搜索关键词: {search_keyword}")
                
                # 处理最终回答
                elif chunk_type == "response.output_text.delta":
                    if not answering_started:
                        answering_started = True
                    delta = getattr(chunk, "delta", "")
                    if delta:
                        print(delta, end="", flush=True)
                        full_response += delta
                        time.sleep(0.01)
                
                # 处理完成事件，提取工具使用统计
                elif chunk_type == "response.completed":
                    if hasattr(chunk, "response") and hasattr(chunk.response, "usage"):
                        usage = chunk.response.usage
                        if hasattr(usage, "tool_usage"):
                            tool_usage = usage.tool_usage
                        if hasattr(usage, "tool_usage_details"):
                            tool_usage_details = usage.tool_usage_details
            
            print()  # 换行
            
            # 提取推理过程和参考资料
            main_response = full_response
            reasoning_process = reasoning_text
            
            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() if reasoning_parts else reasoning_text
            
            # 添加AI回复到对话历史
            self.conversation_history.append({
                "role": "assistant",
                "content": [{"type": "input_text", "text": main_response}]
            })
            
            return {
                "response": main_response,
                "reasoning_process": reasoning_process,
                "references": references,
                "tool_usage": tool_usage,
                "tool_usage_details": tool_usage_details,
                "search_keywords": search_keywords
            }
            
        except Exception as e:
            logger.error(f"对话失败: {str(e)}", exc_info=True)
            return {
                "response": f"抱歉，处理您的请求时出现错误: {str(e)}",
                "reasoning_process": "",
                "references": [],
                "tool_usage": None,
                "tool_usage_details": None,
                "search_keywords": []
            }
    
    def clear_history(self):
        """清除对话历史"""
        self.conversation_history = []
        print("对话历史已清除")
    
    def run(self):
        """运行聊天机器人"""
        print("欢迎使用AI助手！输入 'quit' 退出，输入 'clear' 清除对话历史。")
        if self.enable_web_search:
            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
                
                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(enable_web_search=True)
    bot.run()
