from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models import ChatTongyi
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv
import re
import os

load_dotenv()
api_key = os.getenv('QWEN_API_KEY')

# ============================================================
# 1. 加载 PDF 文档
# ============================================================
loader = PyMuPDFLoader("D:/code/shuheAI/02_RAG/华为擎云W585X 用户指南-(PGUX,KOS&UOS_02,zh-cn).pdf")
pdf_pages = loader.load()

print(f"文档类型：{type(pdf_pages)}")
print(f"PDF 共{len(pdf_pages)}页")

# 查看第一页的内容和元数据
first_page = pdf_pages[0]
print(f"元数据：{first_page.metadata}")
print(f"内容预览：{first_page.page_content[:200]}")

# ============================================================
# 2. 清洗 PDF 文本
# ============================================================
def clean_pdf_text(text: str) -> str:
    """清洗 PDF 解析出的文本，去除常见噪声"""
    # 删除非中文字符之间的换行符
    text = re.sub(r'[^一](\n)[^一]', lambda m: m.group(0).replace('\n', ''), text)
    # 删除项目符号和多余空格
    text = text.replace('•', '').replace('  ', ' ')
    # 删除连续的换行符（保留一个）
    text = re.sub(r'\n{2,}', '\n', text)
    return text.strip()

# 对所有页面清洗
for page in pdf_pages:
    page.page_content = clean_pdf_text(page.page_content)

# ============================================================
# 3. 文本分块
# ============================================================
splitter = RecursiveCharacterTextSplitter(
    separators=["\n\n", "\n", "。", "！", "？", " ", ""],
    chunk_size=50,
    chunk_overlap=10,
    length_function=len
)

split_docs = splitter.split_documents(pdf_pages)
print(f"切分后的文件数量：{len(split_docs)}")
print(f"切分后的字符数（可以用来大致评估 token 数）：{sum([len(doc.page_content) for doc in split_docs])}")

# 过滤掉 page_content 为空或仅含空白的文档
valid_docs = [
    doc for doc in split_docs
    if doc.page_content and doc.page_content.strip()
]
print(f"有效块数量：{len(valid_docs)}")
print(f"总字符数（可大致评估 Token 数）：{sum(len(d.page_content) for d in valid_docs)}")

# ============================================================
# 4. 初始化 Embedding 模型
# ============================================================
embedding_model = DashScopeEmbeddings(
    model="text-embedding-v3",
    dashscope_api_key=api_key
)

# 单条文本向量化
text = "RAG系统搭建实战"
embedding = embedding_model.embed_query(text)
print(f"向量维度：{len(embedding)}")
print(f"前5个值：{embedding[:5]}")

# ============================================================
# 5. 向量化 + 存入向量库
# ============================================================
persist_dir = "./my_knowledge_db"

if os.path.exists(persist_dir) and os.listdir(persist_dir):
    print(f"加载已有向量库：{persist_dir}")
    vectordb = Chroma(
        persist_directory=persist_dir,
        embedding_function=embedding_model
    )
else:
    print(f"创建新向量库：{persist_dir}")
    vectordb = Chroma.from_documents(
        documents=valid_docs,
        embedding=embedding_model,
        collection_metadata={"hnsw:space": "cosine"},  # 余弦相似度
        persist_directory=persist_dir
    )

# 创建检索器，设置返回 Top-3 最相关文档
retriever = vectordb.as_retriever(search_kwargs={"k": 3})

# ============================================================
# 6. 提问
# ============================================================
query = "如何进入BIOS设置？"
relevant_docs = retriever.invoke(query)

for i, doc in enumerate(relevant_docs):
    print(f"--- 结果{i+1} ---")
    print(f"内容：{doc.page_content[:100]}...")
    print(f"来源：{doc.metadata}")
    print()

# 带分数的相似度检索（分数越低越相似，0 表示完全匹配）
results = vectordb.similarity_search_with_score(query, k=3)
for doc, score in results:
    print(f"内容：{doc.page_content[:100]}... | 相似度分数：{score:.4f}")

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

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

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

**上下文：**
{context}

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

llm = ChatTongyi(model="qwen-plus", dashscope_api_key=api_key)
chain = prompt | llm | StrOutputParser()

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