#!/usr/bin/env python3

import os
import sys
from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path
import threading
import time

import numpy as np
import cv2

import rospy # type: ignore
from std_msgs.msg import String # type: ignore
from sensor_msgs.msg import Image as Sensor_Image # type: ignore
from cv_bridge import CvBridge, CvBridgeError # pyright: ignore[reportMissingImports]

#自定义数据结构
from iscas_msgs.msg import LLMDetectCommand

import logging


# 获取当前脚本的路径（Path 对象）
current_script = Path(__file__).resolve()
# 向上两级找到项目根目录（根据实际结构调整层级）
project_root = current_script.parents[1]  # parents[0] 是当前目录，parents[1] 是父目录，以此类推

print("project_root: ",project_root)

from scripts.ai_agent.agents.multimodal_large_model_agent import MMLM_Agent
from tools.core.common_functions import read_yaml_file


# 配置日志记录器
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


# ------------------------------
# ROS1节点封装
# ------------------------------
class MultimodalAgentROSNode:
    """将多模态Agent封装为ROS1节点，提供ROS接口"""
    
    def __init__(self):
        # 初始化ROS节点
        rospy.init_node('multimodal_uav_agent', anonymous=False)
        
        # 从参数服务器获取配置，提供默认值
        # self.vlm_model_path = rospy.get_param(
        #     'uav_agent/vlm_model_path', 
        #     '/path/to/your/llm/model.gguf'
        # )
        # self.embedding_model_path = rospy.get_param(
        #     'uav_agent/embedding_model_path', 
        #     '/path/to/your/embedding/model.gguf'
        # )
        # self.max_context_window = rospy.get_param('uav_agent/max_context_window', 5)
        # self.max_tokens = rospy.get_param('uav_agent/max_tokens', 512)
        # self.persist_directory = rospy.get_param(
        #     'uav_agent/persist_directory', 
        #     os.path.join(os.path.expanduser("~"), ".multimodal_agent", "chroma_db")
        # )

        # # 从参数服务器获取配置，提供默认值
        # self.vlm_model_path = rospy.get_param(
        #     param_name='uav_agent/agent_config_json_file', 
        #     default='/path/to/your/llm/model.gguf'
        # )

        # uav_ai_agent_node_params = read_yaml_file(yaml_path=f"{current_script.parents[0]}/ros_uav_agent_node_param.yaml")
        
        
        # # 验证模型路径
        # if not os.path.exists(self.vlm_model_path):
        #     rospy.logwarn(f"VLM模型路径不存在: {self.vlm_model_path}")
        # if not os.path.exists(self.embedding_model_path):
        #     rospy.logwarn(f"嵌入模型路径不存在: {self.embedding_model_path}")
        
        # 创建持久化目录
        # os.makedirs(self.persist_directory, exist_ok=True)
        # agent_config_json_file = uav_ai_agent_node_params.get("agent_config_json_file")
        agent_config_json_file= f"{current_script.parents[1]}/config/agent_config.json"
        # 初始化Agent
        try:
            self.ai_agent = MMLM_Agent(
                # vlm_model_path=self.vlm_model_path,
                # embedding_model_path=self.embedding_model_path
                # max_context_window=self.max_context_window,
                # max_tokens=self.max_tokens,
                # persist_directory=self.persist_directory
            )
            
            self.ai_agent.parse_config(config_json_file=agent_config_json_file)
            self.ai_agent.init_components(agent_config=self.ai_agent.agent_config)

        except Exception as e:
            rospy.logfatal(f"Agent初始化失败: {str(e)}")
            raise
        
        # 初始化订阅者
        self.text_query_sub = rospy.Subscriber(
            '/uav_agent/text_query',  # 使用私有话题
            String, 
            self.text_query_callback,
            queue_size=1
        )
        
        self.image_sub = rospy.Subscriber(
            '/uav_agent/image_input', 
            Sensor_Image, 
            self.image_callback,
            queue_size=1
        )
        
        self.add_document_sub = rospy.Subscriber(
            '/uav_agent/add_document', 
            String, 
            self.add_document_callback,
            queue_size=1
        )
        
        self.clear_context_sub = rospy.Subscriber(
            '/uav_agent/clear_context', 
            String, 
            self.clear_context_callback,
            queue_size=1
        )

        self.object_detect_sub = rospy.Subscriber(
            '/fastsystem/detect', 
            LLMDetectCommand, 
            self.object_detect_callback,
            queue_size=1
        )
        
        # 初始化发布者
        self.response_pub = rospy.Publisher(
            '/uav_agent/response', 
            String, 
            queue_size=10
        )
        
        self.status_pub = rospy.Publisher(
            '/uav_agent/status', 
            String, 
            queue_size=10
        )
        self.obj_detect_response_pub = rospy.Publisher(
            '/slowsystem/detect_result', 
            String, 
            queue_size=10
        )
        
        # 存储最新的图像
        self.latest_image = None
        self.image_lock = threading.Lock()  # 线程安全锁
        
        # 发布初始化完成状态
        rospy.loginfo("多模态RAG Agent节点初始化完成")
        self.status_pub.publish("状态: 初始化完成，就绪")
        
        # 定时发布状态
        self.status_timer = rospy.Timer(rospy.Duration(10), self.publish_heartbeat)
    
    def publish_heartbeat(self, event):
        """定时发布心跳状态"""
        self.status_pub.publish("状态: 运行中")
    
    def text_query_callback(self, msg: String) -> None:
        """处理文本查询回调"""
        query = msg.data
        if not query.strip():
            rospy.logwarn("收到空的文本查询，忽略")
            self.status_pub.publish("警告: 收到空查询")
            return
            
        rospy.loginfo(f"收到文本查询: {query[:50]}...")  # 只显示前50字符
        self.status_pub.publish(f"处理查询: {query[:30]}...")
        
        # 获取最新图像（如果有）
        image_data = None
        with self.image_lock:
            if self.latest_image is not None:
                image_data = self.latest_image
                self.latest_image = None  # 处理后清除
                rospy.loginfo("使用最新的图像数据进行处理")
        
        # 在单独的线程中处理查询，避免阻塞ROS回调
        threading.Thread(
            target=self._process_query_thread,
            args=(query, image_data),
            daemon=True
        ).start()

    def object_detect_callback(self, msg: LLMDetectCommand) -> None:
        """处理目标检测的回调"""
        target_decription = f"{msg.amount}{msg.property}{msg.target_class}"

        if not target_decription.strip():
            rospy.logwarn("收到空的文本查询，忽略")
            self.status_pub.publish("警告: 收到空查询")
            return
        
        obj_detect_prompt_template= self.ai_agent.prompt_mag.create_object_detect_prompt_template()
        obj_detect_prompt = obj_detect_prompt_template.format(object_description=target_decription)

        # final_prompt_template = self.ai_agent.prompt_mag.create_final_conversation_prompt_template()
        # final_prompt = obj_detect_prompt.format(model_role="AI专家",
        #                                         field="目标检测",
        #                                         context="",
        #                                         chat_history = "",
        #                                         question=obj_detect_prompt)
            
        rospy.loginfo(f"收到文本查询: {obj_detect_prompt[:50]}...")  # 只显示前50字符
        self.status_pub.publish(f"处理查询: {obj_detect_prompt[:30]}...")
        
        # 获取最新图像（如果有）
        #sensor_image 2 dict
        image_data = self.ai_agent.multimodal_processor.ros_image2dict( image_msg=msg.image)  
        # with self.image_lock:
        #     if self.latest_image is not None:
        #         image_data = self.latest_image
        #         self.latest_image = None  # 处理后清除
        #         rospy.loginfo("使用最新的图像数据进行处理")

        obj_detect_query_info = {
            "model_role":"AI专家",
            "field":"图像目标检测",
            "user_query":obj_detect_prompt
        }
        
        query_type = "obj_detect"
        # 在单独的线程中处理查询，避免阻塞ROS回调
        threading.Thread(
            target=self._process_query_thread,
            args=(query_type,obj_detect_query_info,image_data),
            daemon=True
        ).start()
    
    def _process_query_thread(self, query_type:str, query_info: Optional[dict[str,Any]], image_data: Optional[Dict[str, Any]]) -> None:
        """在单独线程中处理查询"""
        try:
            response = self.ai_agent.process_query(query_info = query_info,image_data = image_data)
            # self.response_pub.publish(response)
            if "obj_detect" == query_type:
                self.obj_detect_response_pub.publish(response)
                
                #TODO：
            else:
                self.response_pub.publish(response)
            self.status_pub.publish("查询处理完成")
        except Exception as e:
            error_msg = f"查询处理线程出错: {str(e)}"
            rospy.logerr(error_msg)
            self.response_pub.publish(error_msg)
            self.status_pub.publish(f"错误: {str(e)}")
    
    def image_callback(self, msg: Sensor_Image) -> None:
        """处理图像输入回调"""
        rospy.loginfo(f"收到图像输入，尺寸: {msg.width}x{msg.height}")
        try:
            # 处理图像并存储
            image_data = self.ai_agent.multimodal_processor.ros_image2dict(msg)
            with self.image_lock:
                self.latest_image = image_data
            self.status_pub.publish(f"已接收图像，尺寸: {msg.width}x{msg.height}，等待文本查询")
        except Exception as e:
            rospy.logerr(f"处理图像失败: {str(e)}")
            self.status_pub.publish(f"图像处理失败: {str(e)}")
    
    def add_document_callback(self, msg: String) -> None:
        """添加文档到知识库回调"""
        file_path = msg.data
        rospy.loginfo(f"收到添加文档请求: {file_path}")
        self.status_pub.publish(f"正在处理文档: {os.path.basename(file_path)}")
        
        # 在单独线程中处理文档添加
        threading.Thread(
            target=self._add_document_thread,
            args=(file_path,),
            daemon=True
        ).start()
    
    def _add_document_thread(self, file_path: str) -> None:
        """在单独线程中添加文档"""
        success, message = self.uav_agent.add_knowledge_document(file_path)
        self.status_pub.publish(message)
        rospy.loginfo(f"文档处理结果: {message}")
    
    def clear_context_callback(self, msg: String) -> None:
        """清除上下文回调"""
        rospy.loginfo("收到清除上下文请求")
        self.uav_agent.clear_context()
        self.status_pub.publish("对话上下文已清除")

    # async def uav_conconversation(self):

    #     """
    #     主函数 - 基于LangChain UAV agent对话客户端
    #     """
        
    #     try:

    #         tools = await self.mcp_clients_mag.multi_server_mcp_client.get_tools()
            
    #         # 创建带有智能提示词的代理
    #         prompt_temp = """你是一个专业的智能助手，具备多种工具能力，能够帮助用户完成各类任务。

    #                         可用工具及使用场景：
    #                         - terminal: 执行系统命令、文件操作、程序运行
    #                         - app_analysis: 职业分析、应用程序分析、用户画像  
    #                         - browser: 网页控制、自动化操作、信息采集
    #                         - weather: 天气查询、气象信息获取

    #                         工具使用原则：
    #                         1. 只有用户明确需要执行具体操作时才使用工具
    #                         2. 对于一般性对话、概念解释、建议咨询等，直接回答
    #                         3. 优先选择最合适的工具来完成任务
    #                         4. 如果不确定是否需要工具，优先选择直接回答

    #                         请保持自然、友好的对话风格，准确理解用户真实意图，明智地选择是否使用工具。"""
    #         template_variable_name = "message"
    #         smart_prompt = self.prompt_mag.create_chat_prompt_template(prompt_temp,
    #                                                                                        template_variable_name)
    #         react_agent = create_react_agent(self.model_mag.models_server.model, 
    #                                                        tools, prompt=smart_prompt)

    #         # # 使用窗口记忆 - 保留最近10轮对话，无需transformers包
    #         # memory = ConversationBufferWindowMemory(
    #         #     k=10,                    # 保留最近10轮对话  
    #         #     return_messages=True     # 返回消息对象
    #         # )

    #         print("✨ 全功能智能助手已启动!")
    #         print("🔧 完整工具: 系统命令 | 职业分析 | 浏览器控制 | 天气查询")
    #         print("🧠 AI会智能判断何时需要调用工具")
    #         print("💬 支持流式对话，体验更自然")
    #         print("输入 'exit' 退出\n")

    #         # 对话循环
    #         while True:
    #             user_input = input("💬 问题: ").strip()
                
    #             if user_input.lower() == "exit":
    #                 print("👋 再见!")
    #                 break
                
    #             if not user_input:
    #                 continue
                
    #             # 获取历史消息并添加当前用户输入
    #             messages = memory.chat_memory.messages.copy()
    #             messages.append(HumanMessage(content=user_input))
                
    #             # 先用LLM判断是否需要工具
    #             print("🤔 分析问题中...")
                
    #             # 创建判断提示
    #             judge_prompt = f"""
    #             用户问题: {user_input}
                
    #             请判断这个问题是否需要调用工具来完成任务。
                
    #             需要调用工具的情况(回答YES)：
    #             - 执行系统命令、文件操作、程序运行
    #             - 职业分析、应用程序分析、了解职业匹配度
    #             - 网页控制、信息采集、打开浏览器
    #             - 天气查询、气象信息、天气预报
    #             - 其他需要实际执行操作的任务
                
    #             不需要工具的情况(回答NO)：
    #             - 一般性对话、问候、闲聊
    #             - 概念解释、知识问答
    #             - 建议咨询、意见交流
    #             - 创意写作、故事创作
    #             - 简单数学计算
                
    #             只回答YES或NO，不要其他内容。
    #             """
                
    #             judge_response = await llm.ainvoke([HumanMessage(content=judge_prompt)])
    #             need_tools = "YES" in str(judge_response.content).upper()
                
    #             if need_tools:
    #                 print("🔧 检测到需要工具调用，启动增强模式...")
    #                 ai_response = ""
    #                 tool_used = False
                    
    #                 # 使用agent处理工具调用
    #                 async for chunk in agent.astream({"messages": messages}):
    #                     for node_name, node_output in chunk.items():
    #                         if node_name == "agent":
    #                             messages_in_chunk = node_output.get("messages", [])
    #                             for msg in messages_in_chunk:
    #                                 if msg.type == "ai":
    #                                     if hasattr(msg, "tool_calls") and msg.tool_calls:
    #                                         tool_used = True
    #                                         for tool_call in msg.tool_calls:
    #                                             print(f"🔧 调用工具: {tool_call['name']}")
    #                                             print(f"📝 参数: {tool_call['args']}")
    #                                     elif msg.content:
    #                                         ai_response = msg.content
                            
    #                         elif node_name == "tools":
    #                             messages_in_chunk = node_output.get("messages", [])
    #                             for msg in messages_in_chunk:
    #                                 if msg.type == "tool":
    #                                     print(f"✅ 工具 '{msg.name}' 执行完成")
    #                                     try:
    #                                         import json
    #                                         result = json.loads(msg.content) if isinstance(msg.content, str) else msg.content
    #                                         if isinstance(result, dict):
    #                                             if result.get("success"):
    #                                                 stdout = result.get("stdout", "").strip()
    #                                                 if stdout:
    #                                                     print(f"📄 命令输出:\n{stdout}")
    #                                             else:
    #                                                 stderr = result.get("stderr", "")
    #                                                 print(f"❌ 错误: {stderr}")
    #                                     except:
    #                                         print(f"📄 原始结果: {msg.content}")
                    
    #                 if ai_response:
    #                     print(f"\n🎯 最终回答: {ai_response}")
                        
    #             else:
    #                 print("💭 使用流式回答模式...")
    #                 print("\n🎯 AI回答: ", end="", flush=True)
                    
    #                 # 直接使用LLM流式生成
    #                 ai_response = ""
    #                 async for chunk in llm.astream(messages):
    #                     if chunk.content:
    #                         content_str = str(chunk.content)
    #                         print(content_str, end="", flush=True)
    #                         ai_response += content_str
                    
    #                 print()  # 换行
                
    #             # 保存对话到记忆
    #             if ai_response:
    #                 memory.save_context(
    #                     {"input": user_input},
    #                     {"output": str(ai_response)}
    #                 )
    #                 print("✓ 对话已保存到记忆")
    #             else:
    #                 print("⚠️ 未获取到AI回复")

    #     except Exception as e:
    #         print(f"❌ 初始化失败: {e}")
    #         sys.exit(1)
    #         pass
    
    def run(self) -> None:
        """运行节点"""
        try:
            rospy.spin()
        except KeyboardInterrupt:
            rospy.loginfo("节点被用户中断")
        except Exception as e:
            rospy.logfatal(f"节点运行出错: {str(e)}")
        finally:
            rospy.loginfo("节点关闭")
            self.status_pub.publish("状态: 已关闭")

def main():
    try:
        node = MultimodalAgentROSNode()
        node.run()
    except rospy.ROSInterruptException:
        rospy.logerr("ROS节点初始化失败")
    except Exception as e:
        rospy.logfatal(f"程序致命错误: {str(e)}")



# ------------------------------
# 主函数
# ------------------------------
if __name__ == '__main__':
    main()