import os
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from utils.dataclean import clean_pdf_text
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from config import embedding_model, ChatTongyillm


# ========== 第一步：加载文档 ==========
# 获取当前脚本所在目录，构建 PDF 文件的绝对路径
script_dir = os.path.dirname(os.path.abspath(__file__))
pdf_path = os.path.join(script_dir, "data", "car_info.pdf")

# 创建加载器实例，传入 PDF 文件路径
pdf_loader = PyMuPDFLoader(pdf_path)

# 调用 load() 方法，返回一个 Document 列表（每页一个 Document）
pdf_pages = pdf_loader.load()
print("加载资料完成")


# ========== 第二步：清洗数据（可选，根据文档质量决定）==========
# 清洗每个 Document 的文本内容
for page in pdf_pages:
    page.page_content = clean_pdf_text(page.page_content)
print("清晰资料完成")    


# ========== 第三步：创建递归字符分割器 ==========
text_splitter = RecursiveCharacterTextSplitter(
    # 分隔符优先级：段落 → 换行 → 句号 → 空格 → 硬切
    separators=["\n\n", "\n", "。", "！", "？", " ", ""],
    
    # 每个块最大 50 字符
    chunk_size=50,
    
    # 相邻块重叠 10 字符（chunk_size 的 20%）
    chunk_overlap=10,
    
    # 长度计算函数
    length_function=len
)

# 对整个文档切分
split_docs = text_splitter.split_documents(pdf_pages)
print("切割资料完成") 

# ========== 第四步：向量化 + 存入向量库 ==========
vectorstore = Chroma.from_documents(
    documents=split_docs,
    embedding=embedding_model,
    persist_directory="./knowledge_db"
)
print("存储向量资料完成") 

# ========== 第五步：创建检索器 ==========
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# ========== 第六步：提问 ==========
query = "什么汽车比较经济实惠"
relevant_docs = retriever.invoke(query)

# ========== 第七步：生成回答 ==========
context = "\n\n---\n\n".join([d.page_content for d in relevant_docs])

prompt = ChatPromptTemplate.from_template("""
你是一个专业的知识库助手。请根据以下上下文回答问题。

**规则：**
- 只基于提供的上下文回答，不要编造
- 如果上下文中没有相关信息，直接说「根据现有资料，我找不到这个问题的答案」
- 回答要简洁直接，引用原文时用引号

**上下文：**
{context}

**问题：**
{question}
""")


chain = prompt | ChatTongyillm | StrOutputParser()


if __name__ == "__main__":

    answer = chain.invoke({"context": context, "question": query})
    print(answer)