#!/usr/bin/env python3
import os
import sys
import json
from typing import List, Dict, Any, Optional, Tuple

import magic  # 用于识别文件类型（辅助文件搜索工具）
import base64
import logging
from pathlib import Path

from langchain.agents import initialize_agent, AgentType
from langchain.tools import BaseTool, Tool
from langchain.llms import LlamaCpp
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.vectorstores import Chroma
from langchain.embeddings import LlamaCppEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import (
    TextLoader, UnstructuredFileLoader, PDFMinerLoader,
    UnstructuredMarkdownLoader, Docx2txtLoader
)

from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferWindowMemory
from langchain.schema import HumanMessage, AIMessage
from langchain.output_parsers import PydanticOutputParser 
from langchain_mcp_adapters.client import MultiServerMCPClient


from langgraph.prebuilt import create_react_agent

from llama_cpp import Llama, LlamaGrammar

from requests_toolbelt import ImproperBodyPartContentException

import rospy # type: ignore
import threading
import time
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]
import cv2
import numpy as np

import asyncio
from PIL import Image as PIL_Image

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

print("project_root: ",project_root)

#添加到搜索路径
sys.path.append(str(project_root))

#查看系统环境
print("sys path:",sys.path)

from tools.mcp_clients import MCP_Clients
from tools.core.common_functions import read_json_file
from pnc.pnc import PNC
from models.models_manager import Models_Mag
from agents.memory_manager import MultiModalContextMemory
from agents.prompt_manager import Prompt_Manager
from agents.multimodal_processor import MultimodalProcessor
from agents.rag_manager import RAG_Mag
from models.model_definition import (TextInferenceRequest,MultimodalInferenceRequest,InferenceRequest,InferenceResponse,HealthResponse)

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

class Image_Info:
    type: str
    format: str
    data: str
    width: int
    height: int
    original_encoding: str
    timestamp: float

# ------------------------------
# 主Agent类
# ------------------------------
class MMLM_Agent:
    """多模态 Agent，整合所有核心功能"""
    
    def __init__(self, 
                #  vlm_model_path: str,
                #  embedding_model_path: str
                #  max_context_window: int = 5,
                #  max_tokens: int = 512
                 ):
        """
        初始化多模态UAV Agent
        
        参数:
            vlm_model_path: VLM模型路径
            embedding_model_path: 嵌入模型路径
            # max_context_window: 记忆窗口大小
            # max_tokens: 生成回答的最大token数
        """

        # # 初始化组件
        # self.mcp = MCPProtocol(max_context_length=4096)
        # self.multimodal_processor = MultimodalProcessor()
        # self.rag_processor = RAGProcessor(
        #     embedding_model_path=embedding_model_path,
        #     persist_directory=persist_directory
        # )
        
        # 声明对话记忆
        self.memory_mag = None

        #声明模型管理器
        self.model_mag=None

        #声明mcp_client
        self.mcp_clients_mag = None

        #声明pnc_mag
        self.pnc_mag=None

        #声明prompt_mag
        self.prompt_mag=None

        #agent模态数据处理器
        self.multimodal_processor=None

        #agent的RAG管理工具
        self.rag_mag=None

        #agent配置文件
        self.agent_config = None
        # self.vlm_model_path = vlm_model_path
        # self.embedding_model_path = embedding_model_path
        # self.max_context_window = max_context_window
        # self.max_tokens = max_tokens
        logger.info("多模态UAV Agent初始化完成")

    def parse_config(self,config_json_file):
        """解析agent的config文件"""
        
        self.agent_config= read_json_file(file_path=config_json_file)

        print("finished!")


    def init_memory(self,agent_config):

        """初始化历史记忆的memory"""
        try:
            if "mmcm" == agent_config["memory"]["history_memory_type"]:
                self.memory_mag = MultiModalContextMemory(max_context_window=agent_config["memory"]["max_context_window"],
                                                      history_cnt_lmt=agent_config["memory"]["history_cnt_lmt"],
                                                      is_save_long_history=agent_config["memory"]["is_save_long_history"],
                                                      long_history_cnt_lmt=agent_config["memory"]["long_history_cnt_lmt"],
                                                      long_history_file_path=agent_config["memory"]["long_history_file_path"])
            elif "lccbwm" == agent_config["memory"]["history_memory_type"]:
                self.memory_mag = ConversationBufferWindowMemory(
                    k = agent_config["memory"]["max_context_window"],
                    memory_key=agent_config["memory"]["lccbwm"]["agent_memory_key"],
                    return_messages=agent_config["memory"]["lccbwm"]["return_messages"],
                    output_key=agent_config["memory"]["lccbwm"]["output_key"]
            )

            logger.info(msg=f"memory初始化成功")
            return True

        except Exception as e:
            logger.error(f"memory初始化失败: {str(e)}")
            return False

    def init_model_mag(self,agent_config):
        try:
            # 初始化model mag
            self.model_mag = Models_Mag()

            # if agent_config["model_mag"]["is_model_server"]:
            #     self.model_mag.init_model_server(model_config=agent_config["model_mag"]["model_config"])
            # else:
            #     self.model_mag.init_model_client(model_config=agent_config["model_mag"]["model_config"])
            self.model_mag.init_model_interface(model_config=agent_config["model_mag"]["model_config"],
                                                is_model_server=agent_config["model_mag"]["is_model_server"])

            logger.info(msg=f"model_mag初始化成功")
            return True

        except Exception as e:
            logger.error(f"model_mag初始化失败: {str(e)}")
            return False
        
    def init_mcp_clients_mag(self):
        try:
            # 初始化model mag
            self.mcp_clients_mag = MCP_Clients()
            logger.info(msg=f"mcp_clients_mag初始化成功")
            return True

        except Exception as e:
            logger.error(f"mcp_clients_mag初始化失败: {str(e)}")
            return False
        
    def init_pnc_mag(self):
        try:
            # 初始化pnc_mag
            self.pnc_mag=PNC()
            logger.info(msg=f"pnc_mag初始化成功")
            return True

        except Exception as e:
            logger.error(f"pnc_mag初始化失败: {str(e)}")
            return False
    
    def init_prompt_mag(self):
        try:
            # 初始化prompt_mag
            self.prompt_mag=Prompt_Manager()
            logger.info(msg=f"prompt_mag初始化成功")
            return True

        except Exception as e:
            logger.error(f"prompt_mag初始化失败: {str(e)}")
            return False
        
    def init_multimodal_processor(self):
        try:
            # 初始化multimodal_processor
            self.multimodal_processor=MultimodalProcessor()
            logger.info(msg=f"multimodal_processor初始化成功")
            return True

        except Exception as e:
            logger.error(f"multimodal_processor初始化失败: {str(e)}")
            return False
        
    def init_rag_mag(self,agent_config):
        try:
            # 初始化rag_mag
            self.rag_mag=RAG_Mag()
            self.rag_mag.set_components(agent_config["rag_mag"])
            
            logger.info(msg=f"rag_mag初始化成功")
            return True

        except Exception as e:
            logger.error(f"rag_mag初始化失败: {str(e)}")
            return False
    
    def init_components(self,agent_config):
        try:
            self.init_memory(agent_config=agent_config) 
            self.init_model_mag(agent_config=agent_config) 
            self.init_mcp_clients_mag()
            self.init_pnc_mag() 
            self.init_prompt_mag()
            self.init_multimodal_processor()
            self.init_rag_mag(agent_config=agent_config)
            logger.info(f"agent_components初始化成功")
            return True
        except Exception as e:
            logger.error(f"agent_components初始化失败:{str(e)}")
            return False
        
    # def retrieve_relevant_info(self, query: str, k: int = 3, score_threshold: float = 0.2) -> List[Dict[str, Any]]:
    #     """
    #     检索与查询相关的信息
        
    #     参数:
    #         query: 查询文本
    #         k: 最多返回的结果数量
    #         score_threshold: 相关性分数阈值
            
    #     返回:
    #         包含相关文档内容和分数的列表
    #     """
    #     try:
    #         # 执行相似性搜索，返回带分数的结果
    #         docs_and_scores = self.vectorstore.similarity_search_with_score(query, k=k)
            
    #         # 过滤并格式化结果
    #         results = []
    #         for doc, score in docs_and_scores:
    #             if score < score_threshold:  # 分数越低表示越相似
    #                 results.append({
    #                     "content": doc.page_content,
    #                     "metadata": doc.metadata,
    #                     "similarity_score": float(score)
    #                 })
            
    #         logger.info(f"检索到 {len(results)} 条相关信息 (k={k}, 阈值={score_threshold})")
    #         return results
            
    #     except Exception as e:
    #         logger.error(f"检索相关信息失败: {str(e)}")
    #         return []

    async def process_query(self, query_info: Optional[dict[str,Any]], 
                      image_data: Optional[dict[str,Any]] = None, 
                      max_tokens:int = 300,
                      temperature: float = 0.7, 
                      top_p: float = 0.95,
                      stop: Optional[List[str]] = None,
                      system_prompt = None,
                      is_use_chat_history:bool = False,
                      is_use_rag:bool = False,
                      is_save_history:bool = False) :
        """
        处理用户查询，可包含图像数据
        
        参数:
            query_info: 用户查询的基本信息
            image_data: 可选的图像数据字典
            
        返回:
            生成的回答
        """
        try:
            ## 1.处理文本查询
            processed_text = self.multimodal_processor.process_text(query_info["user_query"])
            # self.mcp.add_message("user", processed_text["data"], "text")
            
            ## 2.处理图像数据（如果有）
            # if image_data:
            #     self.mcp.add_message("user", "图像数据", "image")
            #     logger.info("已添加图像数据到上下文")
            
            ## 3.获取对话历史
            # chat_history = self.memory.load_memory_variables({})["chat_history"]
            # chat_history_str = "\n".join([
            #     f"{'用户' if isinstance(msg, HumanMessage) else '助手'}: {msg.content}"
            #     for msg in chat_history
            # ])
            if is_use_chat_history:
                chat_history = self.memory_mag.get_chat_history()
            # chat_history_str = "\n".join([
            #     f"{'用户' if isinstance(msg, HumanMessage) else '助手'}: {msg.content}"
            #     for msg in chat_history
            # ])

            ## 4.增强检索相关信息
            if is_use_rag:
                relevant_info = self.rag_mag.retrieve_relevant_info(processed_text)
                context_str = "\n\n".join([item["content"] for item in relevant_info])
            
            ## 5.构建提示
            # prompt = self.prompt_template.format(
            #     context=context_str,
            #     question=query,
            #     chat_history=chat_history_str
            # )
    
            # self.prompt_mag.set_conversation_prompt_template(input_variables,conversation_template)
            # "model_role","field","context", "chat_history", "question"
            final_conversation_prompt_template = self.prompt_mag.create_final_conversation_prompt_template()
            final_conversation_prompt = final_conversation_prompt_template.format(
                model_role=query_info["model_role"],
                field=query_info["field"],
                context = context_str,
                chat_history=chat_history,
                question=processed_text
            )
            logger.debug(f"生成提示: {final_conversation_prompt[:200]}...")  # 只显示前200字符
            

            # # 5.是否用system_prompt
            # is_use_system_prompt = False
            # if is_use_system_prompt:
            #     system_prompt = ""
            # else:
            #     system_prompt =  None
            
            ## 6.调用VLM生成回答
            # output = self.llm(
            #     prompt=conversation_prompt,
            #     max_tokens=self.max_tokens,
            #     stop=["\n用户:", "\n助手:", "<|end_of_solution|>"],
            #     temperature=0.7,  # 控制输出随机性
            #     top_p=0.9,
            #     echo=False
            # )
            if image_data: #base64
                multi_modal_request = MultimodalInferenceRequest(user_prompt=final_conversation_prompt,
                                                                 image_data=[image_data["data"]],
                                                                 max_tokens=max_tokens,
                                                                 temperature=temperature,
                                                                 top_p=top_p,
                                                                 system_prompt=system_prompt,
                                                                 stop = stop)
                output = self.model_mag.models_interface.multimodal_inference(request=multi_modal_request)
            else:
                text_request = TextInferenceRequest(user_prompt=final_conversation_prompt,
                                                    max_tokens=max_tokens,
                                                    temperature=temperature,
                                                    top_p=top_p,
                                                    system_prompt=system_prompt,
                                                    stop = stop)
                output=self.model_mag.models_interface.text_inference(request=text_request)
            
            # 提取回答
            response = output["result"]
            
            # 更新记忆和上下文
            if is_save_history:

                if image_data:
                    image_info = {
                        "format": "base64",
                        "data": image_data["data"],
                        "description": "",
                        "timestamp": time.time()
                    }

                    self.memory_mag.add_image_message(image_info)
                # self.memory.save_context(
                #     {"input": processed_text},
                #     {"output": response}
                # )
                human_text_info={
                    "data_type":"text",
                    "format":"human_message",
                    "data": HumanMessage(content=processed_text),
                    "timestamp": time.time()
                }
                self.memory_mag.add_text_message(human_text_info)

                if response:
                    ai_text_info={
                        "data_type":"text",
                        "format":"ai_message",
                        "data": AIMessage(content=response),
                        "timestamp": time.time()
                    }
                    self.memory_mag.add_text_message(ai_text_info)

            # self.mcp.add_message("assistant", response, "text")
            
            logger.info(f"生成回答，长度: {len(response)}")
            return response
            
        except Exception as e:
            error_msg = f"处理查询时出错: {str(e)}"
            logger.error(error_msg)
            # self.mcp.add_message("assistant", error_msg, "text")
            return error_msg 
    
    def add_knowledge_document(self, file_path: str) -> Tuple[bool, str]:
        """添加知识文档到RAG系统"""
        return self.rag_mag.add_document(file_path)
    
    def clear_context(self) -> None:
        """清除对话上下文和记忆"""
        # self.mcp.clear_context()
        self.memory_mag.clear()
        logger.info("已清除所有对话上下文和记忆")
    