
#!/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.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_chroma import Chroma

import threading
import time
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)

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

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

from scripts.ai_agent.tools.memory_mag.rag.rag import (set_embeddings,load_vector_database,
                                                       set_text_splitter,set_document_loaders,
                                                       retrieve_relevant_info)

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

# ------------------------------
# RAG (检索增强生成) 处理器
# ------------------------------
class RAG_Mag:
    """实现检索增强生成功能，管理知识库和检索过程"""
    
    def __init__(self):
        """
        初始化RAG处理器
        
        参数:
            embedding_model_path: 嵌入模型路径
            persist_directory: 向量数据库持久化目录
            chunk_size: 文本分割大小
            chunk_overlap: 文本块重叠大小
        """
        # 初始化嵌入模型
        # self.embedding_model_path = embedding_model_path
        # self.embeddings = LlamaCppEmbeddings(
        #         model_path=embedding_model_path,
        #         n_ctx=n_ctx,
        #         n_threads=n_threads
        #     )
        self.embeddings = None

        # 确保持久化目录存在
        # os.makedirs(persist_directory, exist_ok=True)
            
        # 初始化向量存储
        # self.vectorstore = Chroma(
        #     persist_directory=persist_directory,
        #     embedding_function=self.embeddings,
        #     client_settings={"persist_directory": persist_directory}
        # )
        self.vectorstore = None

        # 初始化文本分割器
        # self.text_splitter = RecursiveCharacterTextSplitter(
        #     chunk_size=chunk_size,
        #     chunk_overlap=chunk_overlap,
        #     length_function=len,
        #     separators=["\n\n", "\n", " ", ""]
        # )
        self.text_splitter = None
        
        # 文档加载器映射
        # self.document_loaders = {
        #     '.txt': TextLoader,
        #     '.pdf': PDFMinerLoader,
        #     '.md': UnstructuredMarkdownLoader,
        #     '.docx': Docx2txtLoader,
        #     # 可以添加更多文档类型
        # }
        self.document_loaders = None

        logger.info(f"RAG_Manager初始化完成")
    
    # def set_embeddings(self,
    #                    embedding_model_name,
    #                    embedding_type:str,
    #                    model_config:Optional[dict[str,Any]]):
    #     """初始化embeddings"""
    #     try:
    #         self.embeddings = set_embeddings(embedding_model_name=embedding_model_name,
    #                                                                             embedding_type=embedding_type,
    #                                                                             model_config=model_config)
    #     except Exception as e:
    #         logger.error(f"embeddings初始化失败:{str(e)}")
    #         return None

    def set_vectorstore(self,
                        embeddings,
                        client_settings:Optional[dict[str,Any]]={},
                        vdb_dir: str = "./chroma_db"
                        ):
        """初始化vectorstore"""
        try:
            # 确保持久化目录存在
            # os.makedirs(persist_directory, exist_ok=True)
            
            # 初始化向量存储
            # 检查是否需要覆盖现有数据库
            if os.path.exists(vdb_dir) :
                print(f"加载现有向量数据库: {vdb_dir}")
                vectorstore = Chroma(
                    embedding_function=embeddings,
                    client_settings=client_settings,
                    persist_directory=vdb_dir
                )
                return vectorstore
            else:
                logger.error(f"vectorstore初始化失败,没有向量数据库")
                return None 
        except Exception as e:
            logger.error(f"vectorstore初始化失败:{str(e)}")
            return None

    # def set_text_splitter(self,
    #                       chunk_size: int = 500,
    #                       chunk_overlap: int = 50):
    #     try:
    #         text_splitter = RecursiveCharacterTextSplitter(
    #             chunk_size=chunk_size,
    #             chunk_overlap=chunk_overlap,
    #             length_function=len,
    #             separators=["\n\n", "\n", " ", ""]
    #         )
    #         return text_splitter
    #     except Exception as e:
    #         logger.error(f"text_splitter初始化失败: {str(e)}")
    #         return None

    # def set_document_loaders(self):
    #     try:
    #         document_loaders = {
    #             '.txt': TextLoader,
    #             '.pdf': PDFMinerLoader,
    #             '.md': UnstructuredMarkdownLoader,
    #             '.docx': Docx2txtLoader,
    #             # 可以添加更多文档类型
    #         }
    #         return document_loaders
    #     except Exception as e:
    #         logger.error(f"document_loaders初始化失败: {str(e)}")
    #         return None
    
    def set_components(self,
                       rag_config:Optional[dict[str,Any]],
                       model_config:Optional[dict[str,Any]]):
        try:
            
            self.embeddings=set_embeddings(embedding_model_path=rag_config["embedding_model_path"],
                                                                                embedding_type=rag_config["embedding_type"],
                                                                                model_config=model_config)
            # self.set_embeddings(embedding_framework_type=rag_config["embedding_framework_type"],
            #            embedding_model_path=rag_config["embedding_model_path"],
            #            n_ctx=rag_config["n_ctx"],
            #            n_threads=rag_config["n_threads"],
            #            n_gpu_layers=rag_config["n_gpu_layers"],
            #            verbose=verbose)
            
            self.vectorstore =load_vector_database(embeddings=self.embeddings,
                                                        vdb_dir=rag_config["vectorstore_persist_directory"],
                                                        collection_name=rag_config["collection_name"]
                                                        # client_settings={"persist_directory": rag_config["vectorstore_persist_directory"]}
                                                        )
            
            self.text_splitter =set_text_splitter(chunk_size=rag_config["text_splitter"]["chunk_size"],
                                   chunk_overlap=rag_config["text_splitter"]["chunk_overlap"])
            
            self.document_loaders=set_document_loaders()

            logger.info(f"set_components成功。")
            return True
        except Exception as e:
            logger.error(f"set_components失败: {str(e)}")
            return False
        
    #创建chromdb向量数据库
    def create_vector_database(embeddings,
                               collection_name,
                               documents,
                               db_dir,
                               overwrite=False):
        """创建或加载向量数据库"""

        # embeddings =  LlamaCppEmbeddings(
        #         model_path = embedding_model_name,
        #         n_threads=n_threads  # 根据 CPU 核心数调整（如 8 核心设为 4 或 8）
        #     )
        # print(f"已加载嵌入模型: {embedding_model_name}")
        try:
            # 检查是否需要覆盖现有数据库
            if os.path.exists(path=db_dir) and not overwrite:
                print(f"加载现有向量数据库: {db_dir}")
                vector_db = Chroma(
                    persist_directory = db_dir,
                    embedding_function=embeddings,
                    collection_name=collection_name
                )
            else:
                print(f"创建新的向量数据库: {db_dir}")
                vector_db = Chroma.from_documents(
                    documents=documents,
                    embedding=embeddings,
                    persist_directory=db_dir,
                    collection_name=collection_name
                )
                vector_db.persist()
            return vector_db
        except Exception as e:
            logger.error(f"create_vector_database失败: {str(e)}")
            return None

    def add_document(self, file_path: str) -> Tuple[bool, str]:
        """
        添加文档到知识库
        
        参数:
            file_path: 文档文件路径
            
        返回:
            (是否成功, 消息)
        """
        if not os.path.exists(file_path):
            error_msg = f"文件不存在: {file_path}"
            logger.error(error_msg)
            return False, error_msg
            
        # 获取文件扩展名
        _, ext = os.path.splitext(file_path.lower())
        if ext not in self.document_loaders:
            error_msg = f"不支持的文件类型: {ext}，支持的类型: {list(self.document_loaders.keys())}"
            logger.error(error_msg)
            return False, error_msg
            
        try:
            # 加载文档
            loader_class = self.document_loaders[ext]
            loader = loader_class(file_path)
            documents = loader.load()
            logger.info(f"加载文档: {file_path}, 页数: {len(documents)}")
            
            # 分割文档
            splits = self.text_splitter.split_documents(documents)
            logger.info(f"文档分割为 {len(splits)} 个片段")

            # 添加到向量存储
            self.vectorstore.add_documents(splits)
            self.vectorstore.persist()
            
            success_msg = f"文档成功添加到知识库: {file_path}，生成 {len(splits)} 个向量"
            logger.info(success_msg)
            return True, success_msg
            
        except Exception as e:
            error_msg = f"添加文档失败: {str(e)}"
            logger.error(error_msg)
            return False, error_msg
    
    def retrieve_relevant_info(self, query: str, k: int = 3, score_threshold: float = 0.2) -> List[Dict[str, Any]]:
        """
        检索与查询相关的信息
        
        参数:
            query: 查询文本
            k: 最多返回的结果数量
            score_threshold: 相关性分数阈值
            
        返回:
            包含相关文档内容和分数的列表
        """
        return retrieve_relevant_info(vectorstore=self.vectorstore,
                                      query=query,
                                      k=k, 
                                      score_threshold=score_threshold)
    
    def get_retriever(self, search_kwargs: Dict[str, Any] = None) -> Any:
        """获取检索器对象，用于LangChain的RetrievalQA链"""
        if search_kwargs is None:
            search_kwargs = {"k": 3}
        return self.vectorstore.as_retriever(search_kwargs=search_kwargs)