# 文件路径：knowledge_qa.py
# pip install langchain langchain-deepseek langchain-chroma chromadb pypdf
from dotenv import load_dotenv
load_dotenv()

import os
from pathlib import Path
from langchain.tools import tool
from langchain.agents import create_agent
from langchain.chat_models import init_chat_model
from langchain.messages import HumanMessage
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader


class KnowledgeBase:
    """个人知识库管理器"""

    def __init__(self, persist_dir: str = "./my_knowledge_db"):
        self.persist_dir = persist_dir
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500, chunk_overlap=50,
            separators=["\n\n", "\n", "。", "！", "？", ". ", "! ", "? ", " "],
        )
        self.vector_store = None
        self._load_or_create()

    def _load_or_create(self):
        """加载已有向量库或创建新的"""
        if os.path.exists(self.persist_dir) and os.listdir(self.persist_dir):
            self.vector_store = Chroma(
                persist_directory=self.persist_dir,
                embedding_function=self.embeddings,
            )
            print(f"已加载向量库：{self.vector_store._collection.count()} 个文档块")
        else:
            self.vector_store = Chroma(
                embedding_function=self.embeddings,
                persist_directory=self.persist_dir,
            )
            print("已创建新的向量库")

    def add_file(self, file_path: str) -> int:
        """添加文件到知识库，返回添加的文档块数"""
        loader = TextLoader(file_path, encoding="utf-8")
        docs = loader.load()

        # 添加文件来源元数据
        for doc in docs:
            doc.metadata["source"] = Path(file_path).name

        chunks = self.text_splitter.split_documents(docs)
        self.vector_store.add_documents(chunks)
        print(f"已添加 {Path(file_path).name}：{len(chunks)} 个文档块")
        return len(chunks)

    def add_text(self, text: str, source: str = "手动添加") -> int:
        """直接添加文本到知识库"""
        chunks = self.text_splitter.create_documents(
            [text], metadatas=[{"source": source}]
        )
        self.vector_store.add_documents(chunks)
        return len(chunks)

    def search(self, query: str, k: int = 3) -> list:
        """搜索知识库"""
        return self.vector_store.similarity_search(query, k=k)

    def get_retriever(self):
        """获取检索器"""
        return self.vector_store.as_retriever(search_kwargs={"k": 3})

# ========== 创建知识库并添加示例数据 ==========

kb = KnowledgeBase("./my_knowledge_db")

# 添加一些示例知识
kb.add_text(
    "菜鸟教程 RUNOOB 的 Python3 基础教程包含以下章节："
    "1. Python 简介与环境搭建 2. 基本数据类型 3. 运算符与表达式 "
    "4. 条件判断 if-else 5. 循环 for/while 6. 函数定义与调用 "
    "7. 模块与包 8. 文件操作 9. 异常处理 10. 面向对象编程",
    source="Python3 教程大纲"
)

kb.add_text(
    "要成为一名优秀的 Python 开发者，建议按以下路线学习："
    "第一步，掌握 Python 基础语法（1-2 周）；"
    "第二步，学习数据结构和算法基础（2-3 周）；"
    "第三步，选择一个方向深入学习（Web 开发/数据分析/AI）；"
    "第四步，做 2-3 个实战项目巩固知识。",
    source="Python 学习路线"
)

kb.add_text(
    "菜鸟教程的在线编程环境支持 Python、JavaScript、Java、C++ 等多种语言。"
    "用户无需安装任何软件，打开浏览器即可编写和运行代码。"
    "在线环境还支持代码高亮、自动补全和错误提示功能。",
    source="在线编程环境说明"
)


# ========== 创建 RAG Agent ==========

@tool
def search_knowledge(query: str) -> str:
    """在个人知识库中搜索相关信息。搜索时使用完整的问题或关键短语。

    Args:
        query: 搜索问题或关键短语
    """
    docs = kb.search(query, k=3)
    if not docs:
        return "知识库中未找到相关信息。"

    results = []
    for i, doc in enumerate(docs, 1):
        source = doc.metadata.get("source", "未知来源")
        content = doc.page_content[:200]
        results.append(f"[{i}] 来源：{source}\n{content}")

    return "\n\n---\n\n".join(results)


model = init_chat_model("deepseek:deepseek-v4-flash", temperature=0)
agent = create_agent(
    model=model,
    tools=[search_knowledge],
    system_prompt="""你是个人知识库助手。

## 规则
1. 所有问题必须先用 search_knowledge 工具检索知识库
2. 回答时注明信息来源（文档名称）
3. 如果知识库中没有相关内容，如实告知
4. 回答要结构化，使用数字列表或分段""",
)


# ========== 测试 ==========

def ask(question: str):
    """提问并流式显示回答"""
    print(f"\n{'='*60}")
    print(f"Q: {question}")
    print(f"{'='*60}")

    result = agent.invoke({
        "messages": [HumanMessage(content=question)]
    })

    # 显示检索到的内容
    for msg in result["messages"]:
        if msg.type == "tool":
            print(f"\n[检索到的内容]")
            print(msg.content[:300])

    print(f"\n[回答]")
    print(result["messages"][-1].content)


ask("Python3 基础教程包含哪些章节？")
ask("如何规划 Python 学习路线？")
ask("菜鸟教程的在线编程环境支持哪些功能？")