import os
import streamlit as st

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.chains import RetrievalQA

# 1. 读取环境变量（在 Hugging Face Space 中添加环境变量 OPENAI_API_KEY）
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

st.set_page_config(page_title="RAG 文档问答 Demo", layout="wide")
st.title("📄 RAG 文档问答 Demo (LangChain + Chroma + Streamlit)")

st.markdown(
    """
    说明：
    1. 上传一份 PDF 文件（例如论文、说明文档）。
    2. 稍等系统构建索引。
    3. 在下方输入问题，将基于文档内容进行回答。
    """
)

# 2. 侧边栏：上传文件
uploaded_file = st.sidebar.file_uploader("上传 PDF 文件", type=["pdf"])

if not OPENAI_API_KEY:
    st.error("请在环境变量中设置 OPENAI_API_KEY 才能调用 OpenAI 接口。")
    st.stop()

if uploaded_file:
    # 把上传的文件临时保存到本地（Space 的临时存储）
    temp_pdf_path = "temp.pdf"
    with open(temp_pdf_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    # 3. 加载 PDF 文档
    loader = PyPDFLoader(temp_pdf_path)
    pages = loader.load()

    # 4. 切分文本
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
    )
    documents = text_splitter.split_documents(pages)

    st.sidebar.write(f"文档页数: {len(pages)}")
    st.sidebar.write(f"切分后的文本块数: {len(documents)}")

    # 5. 构建向量库（Chroma）
    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
    vectorstore = Chroma.from_documents(documents, embedding=embeddings)

    # 6. 构建 RAG QA 链
    llm = ChatOpenAI(
        temperature=0.1,
        model="gpt-4o-mini",  # 或 gpt-4o / gpt-3.5-turbo 等
        openai_api_key=OPENAI_API_KEY
    )

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
        chain_type="stuff",  # 简单拼接检索到的文本块
        return_source_documents=True
    )

    # 7. 用户提问
    user_question = st.text_input("在此输入关于文档的问题：", "")

    if st.button("生成回答") and user_question.strip():
        with st.spinner("正在检索并生成回答，请稍候..."):
            result = qa_chain(user_question)
            answer = result["result"]
            source_docs = result["source_documents"]

        st.subheader("回答：")
        st.write(answer)

        with st.expander("查看参考片段（检索到的文档内容）"):
            for i, doc in enumerate(source_docs):
                st.markdown(f"**片段 {i+1}:**")
                st.write(doc.page_content)
                st.markdown("---")

else:
    st.info("请先在左侧上传一个 PDF 文件。")