import os
from dotenv import load_dotenv
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
# from langchain_community.vectorstores import Chroma
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models import ChatTongyi
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

load_dotenv(dotenv_path=".env")
emb_base_url = os.getenv("EMB_MODEL_URL")
emb_model_name = os.getenv("EMB_MODEL_NAME")
chat_base_url = os.getenv("CHAT_MODEL_URL")
chat_model_name = os.getenv("CHAT_MODEL_NAME")


# ========== 第一步：加载文档 ==========
loader = PyMuPDFLoader("中国近代史.pdf")
pages = loader.load()

# ========== 第二步：清洗数据（可选，根据文档质量决定）==========
# clean_pages = [clean_pdf_text(page) for page in pages]

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



# ========== 第四步：向量化 + 存入向量库 ==========
embedding_model = OpenAIEmbeddings(
    model=emb_model_name,           # 本地模型名称（根据实际部署填写）
    api_key="not-needed",        # 本地服务通常不需要真实 key，填任意值即可
    base_url=emb_base_url  # 本地服务地址，根据实际端口修改
)

# embedding_model = OpenAIEmbeddings(
#     model="BAAI/bge-m3",           # 本地模型名称（根据实际部署填写）
#     api_key="sk-hynxocghcwhevjapdsccungulfggjlaphlxiyenghcotcwux",        # 本地服务通常不需要真实 key，填任意值即可
#     base_url="https://api.siliconflow.cn/v1"  # 本地服务地址，根据实际端口修改
# )

#
# vectorstore = Chroma.from_documents(
#     documents=docs,
#     embedding=embedding_model,
#     collection_metadata={"hnsw:space": "cosine"},
#     persist_directory="./knowledge_db"
# )
#
# exit()
# ========== 加载已有向量库 ==========
vectorstore = Chroma(
    persist_directory="./knowledge_db",  # 之前保存的目录
    embedding_function=embedding_model   # 加载时也需要指定 embedding 模型
)

# ========== 查看向量库前十条chunk ==========
all_data = vectorstore.get(limit=10)

# 查看结构
# print(all_data.keys())
# 输出: dict_keys(['ids', 'embeddings', 'documents', 'metadatas'])

# 查看所有 chunk 的文本内容
for i, doc in enumerate(all_data['documents']):
    print(f"\n--- Chunk {i} ---")
    print(f"ID: {all_data['ids'][i]}")
    print(f"文本: {doc}")
    # print(f"元数据: {all_data['metadatas'][i]}")


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


# # ========== 第六步：提问 ==========
query = "徐中约是谁"

relevant_docs = vectorstore.similarity_search_with_score(query, k=3)

for doc, score in relevant_docs:
    print(f"{doc.page_content}")
    print("-"* 30)


# ========== 第七步：生成回答 ==========
context = "\n\n---\n\n".join([content.page_content for content, score in relevant_docs])
# print(context)
#
prompt = ChatPromptTemplate.from_template("""
你是一个专业的知识库助手。请根据以下上下文回答问题。

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

**上下文：**
{context}

**问题：**
{question}
""")
#
llm = ChatOpenAI(
    model=chat_model_name,
    api_key="api_key",
    base_url=chat_base_url,
    temperature=0.5,
    streaming=True
)
chain = prompt | llm | StrOutputParser()
# print(chain)
#
# answer = chain.invoke({"context": context, "question": query})
for chunk in chain.stream({"context": context, "question": query}):
    print(chunk, end="", flush=True)
# print(answer)