# 该代码用于将本地知识库中的文档导入到ChromaDB中，并使用远程嵌入模型进行向量化
import os
from pathlib import Path
import sys
import chromadb
# from chromadb.utils import embedding_functions - 不再需要
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings, Embeddable
from unstructured.partition.auto import partition
from rich.progress import track
import logging
import requests # 导入requests
import json # 导入json模块
from langchain.schema import Document  # ✅ 新增：用于构造 LangChain 文档对象

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- 配置 ---
# 获取脚本所在目录，确保路径的正确性
SCRIPT_DIR = Path(__file__).resolve().parent
KNOWLEDGE_BASE_DIR = Path("/home/huangfukk/AI_Agent/scripts/ai_agent/memory/knowledge_base/map/src_data")
VECTOR_STORE_DIR = Path("/home/huangfukk/AI_Agent/scripts/ai_agent/memory/knowledge_base/map/vector_store/osm_map1")
COLLECTION_NAME = "drone_docs"
# EMBEDDING_MODEL_NAME = "bge-small-zh-v1.5" # 不再需要，模型名在函数内部处理

# ✅ 新增：定位项目根目录并加入 sys.path，方便复用 rag 模块和工具函数
# 当前文件路径:
#   .../scripts/ai_agent/tools/memory_mag/rag/map_rag/osm_map/ingest.py
# 向上 4 级即可到达 scripts/ai_agent
# parents[0] = map_rag
# parents[1] = rag
# parents[2] = memory_mag
# parents[3] = tools
# parents[4] = ai_agent  ✅
PROJECT_ROOT = SCRIPT_DIR.parents[4]  # /home/.../scripts/ai_agent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.append(str(PROJECT_ROOT))

from tools.core.common_functions import read_json_file
from tools.memory_mag.rag.rag import set_embeddings, build_vector_database

# --- 自定义远程嵌入函数 ---
class RemoteEmbeddingFunction(EmbeddingFunction[Embeddable]):
    """
    一个使用远程、兼容OpenAI API的嵌入服务的嵌入函数。
    """
    def __init__(self, api_url: str):
        self._api_url = api_url
        logging.info(f"自定义嵌入函数已初始化，将连接到: {self._api_url}")

    def __call__(self, input: Embeddable) -> Embeddings:
        """
        对输入的文档进行嵌入。
        """
        # 我们的服务只能处理文本，所以检查输入是否为字符串列表
        if not isinstance(input, list) or not all(isinstance(doc, str) for doc in input):
            logging.error("此嵌入函数仅支持字符串列表（文档）作为输入。")
            return []

        try:
            # 移除 "model" 参数，因为embedding服务可能不需要它
            response = requests.post(
                self._api_url,
                json={"input": input},
                headers={"Content-Type": "application/json"}
            )
            response.raise_for_status()  # 如果请求失败则抛出HTTPError
            
            # 按照OpenAI API的格式解析返回的嵌入向量
            data = response.json().get("data", [])
            if not data:
                raise ValueError("API响应中没有找到'data'字段或'data'为空")

            embeddings = [item['embedding'] for item in data]
            return embeddings

        except requests.RequestException as e:
            logging.error(f"调用嵌入API失败: {e}")
            # 返回一个空列表或根据需要处理错误
            return []
        except (ValueError, KeyError) as e:
            logging.error(f"解析API响应失败: {e}")
            logging.error(f"收到的响应内容: {response.text}")
            return []


def get_documents(directory: Path):
    """从知识库目录加载所有文档并进行切分"""
    documents = []
    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({
                        "text": 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({
                                    "text": 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({
                                "text": 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):
                                    documents.append({
                                        "text": 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({
                            "text": 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 main():
    """主函数，执行文档入库流程"""
    if not KNOWLEDGE_BASE_DIR.exists():
        KNOWLEDGE_BASE_DIR.mkdir(parents=True)
        logging.warning(f"知识库目录不存在，已自动创建: {KNOWLEDGE_BASE_DIR}")
        logging.warning("请向该目录中添加您的知识文件(如 .txt, .pdf, .md)。")
        return

    # 1. 加载并切分文档
    docs_to_ingest = get_documents(KNOWLEDGE_BASE_DIR)
    if not docs_to_ingest:
        logging.warning("在知识库中未找到可处理的文档。")
        return

    # ✅ 2. 读取 rag_config.json，使用与在线 RAG 相同的嵌入配置
    rag_config_path = PROJECT_ROOT / "config" / "rag_config.json"
    logging.info(f"正在读取 RAG 配置: {rag_config_path}")
    rag_config = read_json_file(str(rag_config_path))
    if not rag_config or "rag_mag" not in rag_config:
        logging.error("rag_config.json 中缺少 'rag_mag' 配置，终止入库。")
        return

    rag_mag = rag_config["rag_mag"]

    # 向量库目录与集合名：保持和后端一致
    persist_directory = rag_mag["vectorstore_persist_directory"]
    collection_name = rag_mag.get("collection_name", COLLECTION_NAME)

    # 嵌入模型配置
    embedding_model_path = rag_mag["embedding_model_path"]
    embedding_type = rag_mag.get("embedding_framework_type", "llamacpp_embedding")
    if embedding_type == "llamacpp_embedding":
        embedding_type = "llamacpp_embeddings"  # 与 rag.py 中 set_embeddings 的约定保持一致

    if embedding_type == "llamacpp_embeddings":
        model_config = rag_mag["model_config_llamacpp"]
    else:
        model_config = rag_mag["model_config_huggingFace"]

    logging.info(f"使用嵌入模型: {embedding_model_path}")
    logging.info(f"向量库目录: {persist_directory}, 集合名: {collection_name}")

    # 3. 初始化嵌入模型（与 rag.py 完全一致）
    embeddings = set_embeddings(
        embedding_model_path=embedding_model_path,
        embedding_type=embedding_type,
        model_config=model_config
    )
    if embeddings is None:
        logging.error("嵌入模型初始化失败，终止入库。")
        return

    # 4. 将 docs_to_ingest 转成 LangChain 的 Document 列表
    documents = [
        Document(page_content=doc["text"], metadata=doc["metadata"])
        for doc in docs_to_ingest
    ]
    logging.info(f"准备入库的文档块数量: {len(documents)}")

    # 5. 使用 rag.py 的 build_vector_database 构建/覆盖向量库
    vector_db = build_vector_database(
        documents=documents,
        embeddings=embeddings,
        persist_directory=persist_directory,
        collection_name=collection_name
    )
    if vector_db is None:
        logging.error("构建向量数据库失败。")
        return

    # ✅ 打印统计信息，确认入库成功
    try:
        count = vector_db._collection.count()  # Chroma 内部集合统计
    except Exception:
        count = None
    logging.info(f"数据库中现在有 {count if count is not None else '未知'} 个条目。")

    print("\n✅ 数据入库完成！")
    print(f"知识库位于: {KNOWLEDGE_BASE_DIR}")
    print(f"向量数据库位于: {persist_directory}")


if __name__ == "__main__":
    main() 