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

from langchain.document_loaders import TextLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.document_loaders import (
    TextLoader, UnstructuredFileLoader, PDFMinerLoader,
    UnstructuredMarkdownLoader, Docx2txtLoader
)

from langchain_community.embeddings import HuggingFaceEmbeddings, LlamaCppEmbeddings

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

print("project_root: ",project_root)

# # 从环境变量中获取项目根目录
# project_root = os.environ.get("PROJECT_ROOT")
# if project_root:
#     # 将项目根目录添加到搜索路径
#     sys.path.append(project_root)
#     print(f"已添加项目路径：{project_root}")
# else:
#     raise ValueError("请先设置 PROJECT_ROOT 环境变量")


# 添加到搜索路径,如果该路径不存在则加入
#  str(project_root)
if project_root not in sys.path:
    sys.path.append(str(project_root))


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

#设置huggingFace embedding模型
def set_embeddings_huggingFace(
    model_path: str,
    model_kwargs:Dict[str, Any]={'device': 'cpu'},
    encode_kwargs:Dict[str, Any]={'normalize_embeddings': True}
):
    try:
        # 改用自定义嵌入类，传递底层所需参数
        embeddings = HuggingFaceEmbeddings(
            model_name=model_path,
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs
        )
        return embeddings
    except Exception as e:
        logger.error(f"❌ HuggingFace嵌入模型初始化失败: {str(e)}")
        return None
    
#设置llamacpp embedding模型
def set_embeddings_llamacpp(
    model_path: str,
    n_ctx: int = 512,
    n_threads: int = 4,
    n_gpu_layers: int = 0,
    n_seq_max: int = 128,
    n_threads_batch = 4,
    flash_attn = True,
    verbose: bool = False
):
    try:
        # 改用自定义嵌入类，传递底层所需参数
        embeddings = LlamaCppCustomEmbeddings(
            model_path=model_path,
            n_ctx=n_ctx,
            n_threads=n_threads,
            n_gpu_layers=n_gpu_layers,
            n_seq_max=n_seq_max,
            n_threads_batch = n_threads_batch,
            flash_attn = flash_attn,
            verbose=verbose
        )
        return embeddings
    except Exception as e:
        logger.error(f"❌ llamacpp嵌入模型初始化失败: {str(e)}")
        return None

#设置embeddings模型
def set_embeddings(embedding_model_path,
                   embedding_type:str,
                   model_config:Optional[dict[str,Any]]):

    if "llamacpp_embeddings" == embedding_type:
        embeddings = set_embeddings_llamacpp(
            model_path=embedding_model_path,
            n_ctx=model_config["n_ctx"],
            n_threads=model_config["n_threads"],
            n_gpu_layers=model_config["n_gpu_layers"],  # RTX 5090建议设20，充分利用GPU
            n_seq_max=model_config["n_seq_max"],
            n_threads_batch = model_config["n_threads_batch"],
            flash_attn = model_config["flash_attn"],
            verbose=model_config["verbose"]
        )
    elif "huggingFace_embeddings" == embedding_type:
        embeddings = set_embeddings_huggingFace(
            model_path=embedding_model_path,
            model_kwargs=model_config["model_kwargs"],
            encode_kwargs = model_config["encode_kwargs"]
            )
    else:
        logger.error("unknown embedding type!")
        embeddings = None

    # if embeddings is not None:
    #     test_texts=["测试嵌入","你好"]
    #     test_emb = embeddings.embed_query(test_texts) 
    #     embedding_dim = len(test_emb)
    #     logger.info(f"✅ 自定义嵌入模型加载成功！向量维度：{embedding_dim}")

    return embeddings

#构建向量数据库
def build_vector_database(documents,embeddings,persist_directory,collection_name):

    try:
        os.makedirs(name=persist_directory, 
                    exist_ok=True)
                
        print(f"创建新的向量数据库: {persist_directory}")
        vector_db = Chroma.from_documents(
            documents=documents,
            embedding=embeddings,
            persist_directory=persist_directory,
            collection_name=collection_name
        )
        vector_db.persist()
        print(f"Chroma 已持久化到: {persist_directory}，collection_name={collection_name}")
        return vector_db
    except Exception as e:
        logger.error(f"vector_database构建失败: {str(e)}")
        return None

#加载向量数据库
def load_vector_database(embeddings,
                         persist_directory:str,
                         collection_name:str):
    """加载已存在的向量数据库"""
    try:
        if os.path.exists(path=persist_directory):
            vector_db = Chroma(
                persist_directory=persist_directory,
                embedding_function=embeddings,
                collection_name=collection_name
            )
            print(f"已加载向量数据库 from {persist_directory}")

            return vector_db
        else:
            return None
    except Exception as e:
        logger.error(f"vector_database加载失败: {str(e)}")
        return None

#设置文本分割器   
def set_text_splitter(self,
                          chunk_size: int = 500,
                          chunk_overlap: int = 50,
                          separators: Optional[List]=["\n\n", "\n", " ", ""]):
    try:
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
            separators=separators
        )
        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 retrieve_relevant_info(vectorstore,
                           query: str,
                           k: int = 3, 
                           score_threshold: float = 0.2
                           ) -> List[Dict[str, Any]]:
    """
    检索与查询相关的信息
    
    参数:
        query: 查询文本
        k: 最多返回的结果数量
        score_threshold: 相关性分数阈值
        
    返回:
        包含相关文档内容和分数的列表
    """
    try:
        # 执行相似性搜索，返回带分数的结果
        docs_and_scores = 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 []
    
def get_retriever(vectorstore,search_kwargs: Dict[str, Any] = None) -> Any:
        """获取检索器对象，用于LangChain的RetrievalQA链"""
        if search_kwargs is None:
            search_kwargs = {"k": 3}
        return vectorstore.as_retriever(search_kwargs=search_kwargs)


def test_load_vdb(embedding_model_path,
                  embedding_type,
                  model_config,
                  persist_directory,
                  collection_name):

     #1. 初始化嵌入模型（使用自定义类）
    embeddings = set_embeddings(embedding_model_path=embedding_model_path,
                                    embedding_type=embedding_type,
                                    model_config=model_config)
    
    vector_db= load_vector_database(embeddings=embeddings,
                                         persist_directory=persist_directory,
                                         collection_name=collection_name)
    
    logger.info("load_vdb finished")

    return vector_db

def main():

    persist_directory = "/home/ubuntu/Workspace/Projects/VLM_VLA/UAV_AI/src/Model/AI_Agent/scripts/ai_agent/memory/knowledge_base/map/vector_store/osm_map_4"
    embedding_type = "huggingFace_embeddings"

    model_config_llamacpp={
        "n_ctx":512,
        "n_threads":4,
        "n_gpu_layers":36,
        "n_seq_max":256,
        "n_threads_batch":None,
        "flash_attn":True,
        "verbose":False
    }

    model_config_huggingFace ={
        "model_kwargs": {
            "device": "cuda"
            },

        "encode_kwargs":{
            "normalize_embeddings": True
            }
    }

    if "llamacpp_embeddings" == embedding_type:
        embedding_model_path = "/home/ubuntu/Workspace/Projects/VLM_VLA/Models/Qwen/Qwen3-Embedding-4B-GGUF/Qwen3-Embedding-4B-f16.gguf"
        model_config=model_config_llamacpp

    elif "huggingFace_embeddings" == embedding_type:
        embedding_model_path="/home/ubuntu/Workspace/Projects/VLM_VLA/Models/Qwen/Qwen3-Embedding-4B"
        model_config=model_config_huggingFace
    
    vector_db = test_load_vdb(embedding_model_path=embedding_model_path,
                  embedding_type=embedding_type,
                  model_config=model_config,
                  persist_directory=persist_directory,
                  collection_name="osm_map_docs")
    # query = "学生宿舍10幢"
    # retrieve_result = retrieve_relevant_info(query=query,
    #                        vectorstore=vector_db,
    #                        k=3, 
    #                        score_threshold=0.2)
    # retriever = get_retriever(vectorstore=vector_db,search_kwargs={"k": 3})
    # results = retriever.get_relevant_documents(query)
    try:
        # 创建检索器
        retriever = vector_db.as_retriever(search_kwargs={"k": 3})
        
        # 执行查询
        query = "基础实验楼"
        results = retriever.get_relevant_documents(query)
        
        if results:
            logger.info(f"查询 '{query}' 返回了 {len(results)} 个结果:")
            for i, doc in enumerate(results):
                logger.info(f"\n结果 {i+1}:")
                logger.info(f"内容: {doc.page_content[:100]}...")
                logger.info(f"来源: {doc.metadata.get('source', '未知')}")
        else:
            logger.warning("查询未返回任何结果")
    except Exception as e:
        logger.error(f"查询测试失败: {str(e)}")
    

    print("main() finished!") 

if __name__ == "__main__":
    main()
