# 基于LangChain和Llama.cpp的本地RAG实现
# 安装依赖: pip install langchain llama-cpp-python chromadb pypdf sentence-transformers

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

# from langchain.document_loaders import PyPDFLoader, DirectoryLoader,JSONLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
# from langchain.embeddings import HuggingFaceEmbeddings
# from langchain.vectorstores import Chroma
# from langchain.chains import RetrievalQA
# from langchain.prompts import PromptTemplate

from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.schema import Document
# from langchain_core.embeddings import Embeddings

from langchain_community.document_loaders import PyPDFLoader,DirectoryLoader,JSONLoader
from langchain_community.embeddings import HuggingFaceEmbeddings, LlamaCppEmbeddings

import json # 导入json模块
from unstructured.partition.auto import partition

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

print("project_root: ",project_root)
# print("sys_path: ",sys.path)

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

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

# print("sys_path: ",sys.path)

from tools.memory_mag.rag.llama_cpp_embeddings import LlamaCppCustomEmbeddings
from tools.memory_mag.rag.rag import set_embeddings, build_vector_database
# from llama_cpp import Llama

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

#加载json文件
def load_json_documents(json_dir):
    """加载JSON文件并提取内容"""
    documents = []
    
    # 遍历目录中的所有JSON文件
    for filename in os.listdir(json_dir):
        if filename.endswith('.json'):
            file_path = os.path.join(json_dir, filename)
            
            try:
                # 使用JSONLoader加载文件
                # jq_schema参数用于指定要提取的JSON字段，这里提取所有字段
                loader = JSONLoader(
                    file_path=file_path,
                    jq_schema='.',  # 提取整个JSON内容
                    text_content=False  # 保留JSON结构
                )
                
                # 加载文档
                loaded_docs = loader.load()
                print(f"成功加载JSON文件: {filename}, 包含 {len(loaded_docs)} 个文档对象")
                
                # 为每个文档添加文件名元数据
                for doc in loaded_docs:
                    doc.metadata['source'] = filename
                    documents.append(doc)
                    
            except Exception as e:
                print(f"加载JSON文件 {filename} 时出错: {str(e)}")
    
    if not documents:
        print("警告: 未加载到任何JSON文档，请检查json_dir目录")
    else:
        print(f"总共加载 {len(documents)} 个JSON文档对象")
        
    return documents

#加载ndjson文件
def load_ndjson_documents(file_path: str) -> list[str]:
    """
    加载 NDJSON 文件，提取每行 JSON 中的关键文本字段
    
    Args:
        file_path: NDJSON 文件路径（如 "data/docs.ndjson"）
    
    Returns:
        提取的文本列表（每个元素对应一行 NDJSON 的关键内容）
    """
    documents = []
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"NDJSON 文件不存在：{file_path}")
    
    with open(file_path, "r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:  # 跳过空行
                continue
            try:
                # 解析每行 JSON（需根据你的 NDJSON 结构调整字段名）
                json_obj = json.loads(line)
                # 示例：提取 "content" 字段（若你的字段是 "text"，则改为 json_obj["text"]）
                if "content" not in json_obj:
                    print(f"警告：第 {line_num} 行 NDJSON 缺少 'content' 字段，跳过")
                    continue
                text = json_obj["content"].strip()
                if text:  # 跳过空文本
                    documents.append(text)
            except json.JSONDecodeError as e:
                print(f"警告：第 {line_num} 行 NDJSON 解析失败：{e}，跳过")
    return documents

#加载文档
def load_documents(file_mame: str):
    """从知识库目录加载所有文档并进行切分"""

    documents = []
    file_path = Path(file_mame)
    # logging.info(f"从 '{directory}' 加载文档...")
    # for file_path in directory.rglob("*"):
    if file_path.is_file() and not file_path.name.startswith('.'):
        try:
            # 对简单文本文件直接读取
            if file_path.suffix in ['.txt', '.md']:
                text = file_path.read_text(encoding='utf-8')
                documents.append(Document(
                    page_content=text,
                    metadata={"source": str(file_path.name)}
                ))
                logging.info(f"成功处理文本文件: {file_path.name}")
            # 特别处理常规的JSON文件
            elif file_path.suffix == '.json':
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    if 'elements' in data and isinstance(data['elements'], list):
                        for element in data['elements']:
                            documents.append(Document(
                                page_content=json.dumps(element, ensure_ascii=False),
                                metadata={"source": str(file_path.name)}
                            ))
                        logging.info(f"成功处理JSON文件: {file_path.name}, 提取了 {len(data['elements'])} 个元素。")
                    else:
                        documents.append(Document(
                            page_content=json.dumps(data, ensure_ascii=False),
                            metadata={"source": str(file_path.name)}
                        ))
                        logging.info(f"成功处理JSON文件: {file_path.name} (作为单个文档)")
            # 新增：专门处理我们生成的 NDJSON 文件
            elif file_path.suffix == '.ndjson':
                with open(file_path, 'r', encoding='utf-8') as f:
                    count = 0
                    for line in f:
                        try:
                            record = json.loads(line)
                            if 'text' in record and isinstance(record['text'], str):

                                # print(f"src text: {record['text']}")
                                record['text']=text_clean_process(record['text'])
                                # print(f"processed text: {record['text']}")

                                # logger.info(f"该句文本token数量:{len(record['text'])}")
                                documents.append(Document(
                                    page_content=record['text'],
                                    metadata={"source": str(file_path.name)}
                                ))
                                count += 1
                        except json.JSONDecodeError:
                            logging.warning(f"跳过无效的JSON行: {line.strip()}")
                if count > 0:
                    logging.info(f"成功处理NDJSON文件: {file_path.name}, 提取了 {count} 个文档。")
            # 对其他所有文件类型，使用unstructured
            else:
                elements = partition(filename=str(file_path))
                for element in elements:
                    documents.append(Document(
                        page_content=element.text,
                        metadata={"source": str(file_path.name)}
                    ))
                logging.info(f"成功处理文件: {file_path.name} (使用unstructured)")
        except Exception as e:
            logging.error(f"处理文件 {file_path.name} 失败: {e}")
    return documents

#内容分块
def process_and_split_documents(self, documents):
    """处理并分割文档"""
    # 对于JSON数据，我们需要将其转换为字符串并适当分割
    # 首先将每个JSON文档转换为格式化的字符串
    for doc in documents:
        # 如果内容是字典，转换为格式化的JSON字符串
        if isinstance(doc.page_content, dict):
            doc.page_content = json.dumps(doc.page_content, indent=2, ensure_ascii=False)
    
    # 创建文本分割器，考虑JSON结构特点设置分割参数
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=100,
        separators=["}\n", "},\n", "\n\n", "\n", " ", ""]  # 优先按JSON结构分割
    )
    
    chunks = text_splitter.split_documents(documents)
    print(f"文档分割完成，得到 {len(chunks)} 个片段")
    
    return chunks

#文本清洗
def text_clean_process(text:str):
    # 增强文本清洗
    text = text.strip().replace("\n", " ").replace("\r", "").replace("\t", " ")
    # 移除多余空格和非UTF-8字符
    text = " ".join(text.split())
    text = text.encode("utf-8", errors="ignore").decode("utf-8")
    return text


def build_osm_map_knowledge_base(
    osm_map_file,
    embedding_model_path,
    embedding_type:str,
    model_config:Optional[dict[str,Any]],
    persist_directory,
    collection_name =None
    ):
    """创建osm_map向量数据库（确保嵌入函数有效）"""
    ## 1. 加载文档（已处理文本长度，单条≈100token < 512）
    docs_to_ingest = load_documents(osm_map_file)
    if not docs_to_ingest:
        logging.warning("⚠️ 在知识库中未找到可处理的文档。")
        return None

    ## 2. 初始化嵌入模型（使用自定义类）
    embeddings = set_embeddings(embedding_model_path=embedding_model_path,
                                    embedding_type=embedding_type,
                                    model_config=model_config)
    # 关键：检查嵌入模型是否初始化成功
    if embeddings is None:
        logger.error("❌ 嵌入模型未初始化，终止向量库构建")
        return None
    return build_vector_database(
                           documents=docs_to_ingest, 
                           persist_directory=persist_directory, 
                           embeddings=embeddings,
                           collection_name=collection_name)

if __name__ == "__main__":
    osm_map_file = "/home/ubuntu/Workspace/Projects/VLM_VLA/UAV_AI/src/Model/AI_Agent/scripts/ai_agent/memory/knowledge_base/map/src_data/export.ndjson"
    
    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"
    # n_threads = 8  # 建议设为CPU核心数的50%-80%（如16核设8-12）
    # 调整GPU层数为20，加速模型加载和嵌入生成
    # "llamacpp_embeddings" /"huggingFace_embeddings"
    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
        
    build_osm_map_knowledge_base(
        osm_map_file=osm_map_file,
        embedding_model_path=embedding_model_path,
        embedding_type=embedding_type,
        model_config=model_config,
        persist_directory=persist_directory,
        collection_name="osm_map_docs")
