# 文件路径：knowledge_qa.py
# pip install langchain langchain-deepseek langchain-openai langchain-community 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, PyPDFLoader


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

    def __init__(self, persist_dir: str = "./my_knowledge_db"):
        self.persist_dir = persist_dir

        # 阿里云百炼的 Embedding 接口兼容 OpenAI 规范，用 langchain-openai
        # 的 OpenAIEmbeddings 调用即可，不需要装 langchain-community 里
        # 已停止维护的 DashScopeEmbeddings。
        # check_embedding_ctx_length=False：关掉 tiktoken 预分词，直接发原始文本
        # （百炼接口不接受 token id 数组）。
        # chunk_size=10：百炼 Embedding 接口单次请求最多接受 10 条文本。
        self.embeddings = OpenAIEmbeddings(
            model="text-embedding-v4",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            check_embedding_ctx_length=False,
            chunk_size=10,
        )
        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):
        """加载已有向量库或创建新的（Chroma 用同一套参数即可，
        目录是否已存在只影响提示信息，不影响创建逻辑）"""
        is_existing = os.path.exists(self.persist_dir) and os.listdir(self.persist_dir)

        self.vector_store = Chroma(
            embedding_function=self.embeddings,
            persist_directory=self.persist_dir,
        )

        if is_existing:
            print(f"已加载向量库：{self.vector_store._collection.count()} 个文档块")
        else:
            print("已创建新的向量库")

    def add_file(self, file_path: str) -> int:
        """添加文件到知识库，返回添加的文档块数

        根据文件扩展名自动选择合适的 Loader：
        - .pdf 用 PyPDFLoader（依赖 pypdf 库解析 PDF，按页返回多个 Document）
        - 其余（.md、.txt 等纯文本文件）用 TextLoader 按原始文本读取
        """
        suffix = Path(file_path).suffix.lower()
        if suffix == ".pdf":
            loader = PyPDFLoader(file_path)
        else:
            loader = TextLoader(file_path, encoding="utf-8")

        docs = loader.load()

        # 统一补充文件名来源；PyPDFLoader 加载出的每个 Document 还会自带
        # page 字段（第几页），检索时可以一并用来定位片段位置
        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="在线编程环境说明"
)

# 也可以加载本地文件，PDF 和 Markdown/TXT 都会自动识别：
# kb.add_file("./docs/产品手册.pdf")
# kb.add_file("./docs/常见问题.md")


# ========== 创建 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", "未知来源")
        page = doc.metadata.get("page")
        location = f"{source}" + (f" 第 {page + 1} 页" if page is not None else "")
        content = doc.page_content[:200]
        results.append(f"[{i}] 来源：{location}\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. 回答时注明信息来源（文档名称，如果是 PDF 还要注明页码）
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)


# ========== 测试：流式，对应系统设计里的"流式输出" ==========

def ask_stream(question: str):
    """提问并逐 Token 流式打印回答，实现打字机效果"""
    print(f"\n{'='*60}")
    print(f"Q: {question}")
    print(f"{'='*60}")
    print("\n[回答] ", end="", flush=True)

    for chunk, metadata in agent.stream(
        {"messages": [HumanMessage(content=question)]},
        stream_mode="messages",
    ):
        # metadata["langgraph_node"] == "model" 表示这个 chunk 来自模型生成
        # 最终回答的节点，过滤掉工具调用等其他类型的 chunk
        if metadata.get("langgraph_node") == "model" and chunk.content:
            print(chunk.content, end="", flush=True)
    print()


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