#!/usr/bin/env python3

import os
import sys
from typing import List, Dict, Any, Optional, Tuple, Union

import base64
import json
import logging
from pathlib import Path
import time

from langchain.memory import ConversationBufferWindowMemory
from langchain.prompts import PromptTemplate
from langchain.schema import HumanMessage, AIMessage

from langchain_core.memory import BaseMemory
from langchain.schema import Document, BaseMessage, HumanMessage, AIMessage

import cv2
from std_msgs.msg import String
from sensor_msgs.msg import Image as SensorImage
from cv_bridge import CvBridge
from PIL import 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))

from tools.core.common_functions import persist_listdata2ndjson,persist_list2json

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

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

#带有聊天记录分类的历史记忆
class Memory_Store:
    def  __init__(self):
        self.memory={}
    
    def get_history(self,context_id):
        return self.memory.get(context_id,[])
    
    def append(self,context_id, user_input, ai_response):
        if context_id not in self.memory:
            self.memory[context_id]=[]
        self.memory[context_id].append({"user": user_input,"ai":ai_response})




# 上下文工程实现 - 支持多模态
# class MultiModalContextMemory(BaseMemory):
class MultiModalContextMemory():
    def __init__(self,
                 max_context_window: int = 4000,
                 history_cnt_lmt:int = 8 ,
                 is_save_long_history:bool=False,
                 long_history_cnt_lmt:int = 100,
                 long_history_file_path:str=""):
        """初始化多模态上下文管理器"""
        self.max_context_window = max_context_window
        self.chat_history: List[Union[BaseMessage, Dict[str, Any]]] = []
        self.bridge = CvBridge()
        self.history_cnt_lmt = history_cnt_lmt
        # 长上下文记忆
        self.is_save_long_history= is_save_long_history
        self.long_history:List[Union[BaseMessage, Dict[str, Any]]] = []
        self.long_history_cnt_lmt:int = long_history_cnt_lmt
        # self.long_history_cnt:int = 0
        self.long_history_file_path = long_history_file_path
    
    def add_text_message(self, 
                         message: Optional[dict[str,Any]],
                         is_save_long_history:bool = False) -> None:
        """
        添加文本消息到历史记录
        
        src_text_info: 文本数据信息集合 类型为dict
            包含{
                "data_type": 数据类型，"text",
                "format": 数据格式，"ai_message",
                "data": 数据， AIMessage(content=response),
                "timestamp": 时间戳， time.time()
            }
        """
        # src_text_info={
        #         "data": message,
        #         "timestamp": 
        #     }

        self.chat_history.append(message)
        self._truncate_history(self.history_cnt_lmt)

        if is_save_long_history:
            self.long_history.append(message)
            # self.long_history_cnt = self.long_history_cnt+1
            long_history_len = len(self.long_history)
            if long_history_len >= self.long_history_cnt_lmt:
                # 保存数据到硬盘
                persist_file_path = f"{self.long_history_file_path}_{self.long_history[0]["timestamp"]}_{self.long_history[long_history_len-1]["timestamp"]}.json"
                persist_list2json(data_list=self.long_history,
                                        file_path=persist_file_path)
                self.long_history.clear()
                # self.long_history_cnt=0
    
    def add_image_message(self,
                          src_image_info:Optional[dict[str,Any]],
                          is_save_long_history:bool = False
                          ) -> None:
        """
        添加图像消息到历史记录

        params:
            src_image_info: 图像数据信息集合 类型为dict
            包含{
                "data_type": 数据类型
                "format": 图像的格式
                "data": 图像数据
                "description": 图像描述文本
                "timestamp": 图像的时间戳
            }
        """
        # 转换为PIL图像以便后续处理
        try:
            # cv_image = self.bridge.imgmsg_to_cv2(src_image, desired_encoding='bgr8')
            # pil_image = Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
            
            # self.chat_history.append({
            #     "type": "image",
            #     "format":src_image_info["format"],
            #     "content": src_image_info["data"],
            #     "description": src_image_info["description"],
            #     "timestamp": src_image_info["timestamp"]
            # })
            self.chat_history.append(src_image_info)
            self._truncate_history(self.history_cnt_lmt)
            logger.info("图像已添加到对话历史")

            if is_save_long_history:
                self.long_history.append(src_image_info)
                # self.long_history_cnt = self.long_history_cnt+1
                long_history_len = len(self.long_history)
                if long_history_len >= self.long_history_cnt_lmt:
                    # 保存数据到硬盘
                    persist_file_path = f"{self.long_history_file_path}_{self.long_history[0]["timestamp"]}_{self.long_history[long_history_len-1]["timestamp"]}.json"
                    persist_list2json(data_list=self.long_history,
                                    file_path=persist_file_path)
                    self.long_history.clear()
        except Exception as e:
            logger.error(f"添加图像到历史记录失败: {str(e)}")
    
    def _truncate_history(self,history_cnt_lmt) -> None:
        """截断历史记录以满足最大token限制"""
        # 简单实现：限制最多history_cnt_lmt轮对话/图像
        if len(self.chat_history) > history_cnt_lmt:
            if history_cnt_lmt > 0:
                self.chat_history = self.chat_history[-history_cnt_lmt:]  
            else:
                self.chat_history = []
                return
    
    def get_k_chat_history(self, k):
        """获取k个上文对话历史"""
        E_idx = len(self.chat_history) - 1
        S_idx = E_idx + 1 - k if E_idx + 1 > k else 0
        context = "以下是对话历史：\n"
        for i in range(S_idx,E_idx) :
            item = self.chat_history[i]
            # if isinstance(item, HumanMessage):
            if isinstance(item, dict) \
                and item["data_type"]=="text" \
                and item["format"]=="human_message":

                context += f"用户: {item.content}\n"
            # elif isinstance(item, AIMessage):
            elif isinstance(item, dict) \
                and item["data_type"]=="text" \
                and item["format"]=="ai_message":
                
                context += f"助手: {item.content}\n"
            elif isinstance(item, dict) \
                and item["data_type"] == "image":
                
                context += f"用户发送了一张图像: {item['description']}\n"
        return context
    
    def get_k_context_query(self, query: str,k) -> str:
        """构建k个包含历史对话和当前查询的完整上下文"""
        context =  self.get_k_chat_history(k)
        context += f"\n当前查询: {query}\n"
        return context
    
    def get_chat_history(self) -> str:
        """获取历史上文对话"""
        context = "以下是对话历史：\n"
        for item in self.chat_history:
            # if isinstance(item, HumanMessage):
            if isinstance(item, dict) \
                and item["data_type"]=="text" \
                and item["format"]=="human_message":
                
                context += f"用户: {item.content}\n"
            # elif isinstance(item, AIMessage):
            if isinstance(item, dict) \
                and item["data_type"]=="text" \
                and item["format"]=="ai_message":

                context += f"助手: {item.content}\n"
            elif isinstance(item, dict) \
                and item["type"] == "image":

                context += f"用户发送了一张图像: {item['description']}\n"
        return context
    
    def get_context_query(self, query: str) -> str:
        """构建包含历史对话和当前查询的完整上下文"""
        context = self.get_chat_history()
        context += f"\n当前查询: {query}\n"
        return context
    
    def get_recent_images(self, k: int = 2) -> List[Image.Image]:
        """获取最近的k张图像"""
        images = []
        for item in reversed(self.chat_history):
            if isinstance(item, dict) and item["type"] == "image":
                images.append(item["content"])
                if len(images) >= k:
                    break
        return list(reversed(images))  # 保持时间顺序
    
    def clear(self) -> None:
        """清空对话历史"""
        self.chat_history = []
        logger.info("多模态对话上下文已清空")

    # def save_context(self, inputs: dict[str, Any], outputs: dict[str, str]) -> None:
    #     """Save context from this conversation to buffer."""
    #     input_str, output_str = self._get_input_output(inputs, outputs)
    #     self.chat_history.add_messages(
    #         [
    #             HumanMessage(content=input_str),
    #             AIMessage(content=output_str),
    #         ],
    #     )


    # async def asave_context(
    #     self,
    #     inputs: dict[str, Any],
    #     outputs: dict[str, str],
    # ) -> None:
    #     """Save context from this conversation to buffer."""
    #     input_str, output_str = self._get_input_output(inputs, outputs)
    #     await self.chat_memory.aadd_messages(
    #         [
    #             HumanMessage(content=input_str),
    #             AIMessage(content=output_str),
    #         ],
    #     )

    # def clear(self) -> None:
    #     """Clear memory contents."""
    #     self.chat_memory.clear()

    # async def aclear(self) -> None:
    #     """Clear memory contents."""
    #     await self.chat_memory.aclear()

    def persist_context_ndjson(self,
                        persist_directory:str,
                        rw_mode:str="w",
                        encoding:str="utf-8",
                        start:Optional[int] = None,
                        end:Optional[int] = None):
        
        chat_history_len = len(self.chat_history)
        if start is not None:
            S_idx = start if start > 0 and start < chat_history_len else 0
        else:
            S_idx = 0
        
        
        if end is not None:
            E_idx = end if end < chat_history_len  else chat_history_len - 1
        else:
            E_idx = chat_history_len - 1

        #保存到NDJSON文件
        with open(file=persist_directory, mode=rw_mode, encoding=encoding) as f:
            for i in range(S_idx,E_idx) :
                item = self.chat_history[i]

                if isinstance(item["data"], HumanMessage):
                    context_line={"用户": item.content}
                    # context += f"用户: {item.content}\n"
                elif isinstance(item["data"], AIMessage):
                    context_line={"助手": item.content}
                    # context += f"助手: {item.content}\n"
                elif isinstance(item, dict) and item["type"] == "image":
                    # context += f"用户发送了一张图像: {item['description']}\n"
                    context_line={"用户": f"用户发送了一张图像: {item['description']}"}
                # context_ndjson.append(context_line)
                json_str = json.dumps(context_line, ensure_ascii=False)

                f.write(json_str + "\n")


    def load_context_ndjson(self,
                            file_path:str,
                            rw_mode:str="r",
                            encoding:str="utf-8",
                            start:Optional[int] = None,
                            end:Optional[int] = None):
        """
        读取NDJSON文件，返回解析后的字典列表（每行一个字典）
        
        :param file_path: NDJSON文件路径（字符串或Path对象）
        :return: 包含所有JSON对象的列表（每个元素为dict）
        """
        data = []
        file_path = Path(file_path)  # 统一转为Path对象，方便处理路径
    
        # 检查文件是否存在
        if not file_path.exists():
            print(f"警告：文件不存在 - {file_path}")
            return data
        
        # 检查是否是文件（不是目录）
        if not file_path.is_file():
            print(f"警告：{file_path} 不是文件")
            return data
        
        # 逐行读取并解析
        with open(file_path, mode=rw_mode, encoding=encoding) as f:
            line_count = len(f.readlines())
            if end is not None:

                E_idx = end if end < line_count-1 else line_count-1
            else:
                E_idx = line_count-1

            # line_count = len(f.readlines())
            if start is not None:
                S_idx = start if start > max(E_idx - self.history_cnt_lmt + 1, 0) else max(E_idx - self.history_cnt_lmt + 1, 0)
            else:
                S_idx = max(E_idx - self.history_cnt_lmt + 1, 0)

            for line_num, line in enumerate(f, start=0):  # 记录行号，方便定位错误
                if line_num >= S_idx and line_num < E_idx:
                    line = line.strip()  # 去除首尾空白（如换行符、空格）
                    
                    # 跳过空行（避免解析空字符串报错）
                    if not line:
                        continue
                    
                    try:
                        # 将JSON字符串解析为字典，添加到列表
                        json_obj = json.loads(line)
                        data.append(json_obj)
                    except json.JSONDecodeError as e:
                        # 捕获无效JSON格式的错误，提示行号便于排查
                        print(f"警告：第 {line_num} 行JSON格式无效，已跳过 - 错误：{e}")
    
        return data
    
    def load_context_json(self,
                          file_path:str,
                          rw_mode:str="r",
                          encoding:str="utf-8",
                          start:Optional[int] = None,
                          end:Optional[int] = None):
        """
        读取NDJSON文件，返回解析后的字典列表（每行一个字典）
        
        :param file_path: NDJSON文件路径（字符串或Path对象）
        :return: 包含所有JSON对象的列表（每个元素为dict）
        """
        data = []
        file_path = Path(file_path)  # 统一转为Path对象，方便处理路径
    
        # 检查文件是否存在
        if not file_path.exists():
            print(f"警告：文件不存在 - {file_path}")
            return data
        
        # 检查是否是文件（不是目录）
        if not file_path.is_file():
            print(f"警告：{file_path} 不是文件")
            return data
        
        try:
            # 逐行读取并解析
            with open(file_path, mode=rw_mode, encoding=encoding) as f:
                src_data = json.load(f)
            
            src_data_len = len(src_data)

            if end is not None:

                E_idx = end if end < src_data_len-1 else src_data_len-1
            else:
                E_idx = src_data_len-1

            # line_count = len(f.readlines())
            if start is not None:
                S_idx = start if start > max(E_idx - self.history_cnt_lmt + 1, 0) else max(E_idx - self.history_cnt_lmt + 1, 0)
            else:
                S_idx = max(E_idx - self.history_cnt_lmt + 1, 0)

            if isinstance(src_data, dict):
                data.append(src_data)
                print("JSON是单个dict，已添加到list")
            
            # 情况2：JSON内容是list → 检查元素是否为dict，过滤非dict元素
            elif isinstance(src_data, list):
                # 遍历列表，只保留dict类型的元素
                for item in data:
                    if isinstance(item, dict):
                        data.append(item)
                    else:
                        print(f"警告：列表中存在非dict元素 {item}，已跳过")
                print(f"JSON是list，共保留 {len(src_data)} 个dict元素")
            
            # 其他情况（如JSON是字符串、数字等，无法转为dict列表）
            else:
                print(f"错误：JSON内容类型为 {type(data)}，无法提取dict对象")
    
        except json.JSONDecodeError as e:
            print(f"错误：JSON格式无效 - {e}（文件：{file_path}）")
        except Exception as e:
            print(f"加载文件失败：{e}（文件：{file_path}）")
    
        return data
    





    