import streamlit as st
from langchain.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import ConversationalRetrievalChain
from langchain.llms import OpenAI
import os
import tempfile
from doc_qa import  embeddings,llm
from doc_qa_1 import embeddings,doc_qa

def start_message(doc_name):
    st.success("✅ ドキュメントのアップロードが完了しました！")
    st.markdown(f"### 📄 アップロードされました: `{doc_name}`")
    st.markdown("これで文書に関する質問ができます。 💬")
    st.markdown("例えば、次のような質問ができます。:")
    st.markdown("- この文書は何について書かれていますか？")
    st.markdown("- 重要なポイントを要約してください。")
    st.markdown("- 著者は誰ですか？")
    st.markdown("はじめるには、下に質問を入力してください。!")

# Function to load individual file
def load_file(file, suffix):
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
        temp_file.write(file.read())
        temp_file_path = temp_file.name

    if suffix == ".pdf":
        loader = PyPDFLoader(temp_file_path)
    elif suffix == ".docx":
        loader = Docx2txtLoader(temp_file_path)
    elif suffix == ".txt":
        loader = TextLoader(temp_file_path)
    else:
        return []

    return loader.load()
st.set_page_config(
    page_title="QA Assistant",
    page_icon="https://yourdomain.com/logo.png",
    layout="centered"
)
# Title
st.title("📄 ドキュメント質問応答支援ツール")

# Step 1: Upload document
if "file_uploaded" not in st.session_state:
    st.session_state.file_uploaded = False
    st.markdown("""
    👋 こちらへようこそ！私は文書の内容を理解するためのインテリジェントアシスタントです。

あなたは以下のことができます：

PDF、DOCX、TXTファイルをアップロード

文書の内容について質問

要約、重要ポイント、または具体的な詳細の取得

🛠️ 質問の例：
この文書は何について書かれていますか？

主要なポイントを要約してください。

著者は誰ですか？

重要な日付や締め切りは何ですか？

結論や推奨事項は何ですか？

📂 まず、1つ以上の文書をアップロードしてください。
💬 その後、下に質問を入力しましょう！
    """)
if "messages" not in st.session_state:
    st.session_state.messages = []


flag = 0
# Upload multiple files
with st.sidebar:
    uploaded_files = st.file_uploader("PDF、DOCX、またはTXTファイルをアップロードしてください。", type=["pdf", "docx", "txt"], accept_multiple_files=True)
    # Load and process documents
    file_names=[]
    if uploaded_files:
        all_docs = []
        for file in uploaded_files:
            suffix = os.path.splitext(file.name)[1]
            docs = load_file(file, suffix)
            all_docs.extend(docs)
            file_names.append(file.name)

        # Split and embed documents
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
        split_docs = text_splitter.split_documents(all_docs)
        #embeddings = OpenAIEmbeddings()
        vectorstore = FAISS.from_documents(split_docs, embeddings)

        # Setup ConversationalRetrievalChain
        qa_chain = ConversationalRetrievalChain.from_llm(
            llm=llm,
            retriever=vectorstore.as_retriever(),
            return_source_documents=False
        )
        start_message('\n'.join(file_names))
        flag = 1

# Initialize session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

if flag==1:
    if user_query := st.chat_input():
        st.session_state.messages.append({"role": "user", "content": user_query})
        with st.chat_message("user"):
            st.markdown(user_query)
            #st.markdown(f"**Q:** {user_query}")
        result=doc_qa(user_query,vectorstore)
        st.session_state.messages.append({"role": "assistant", "content": result["answer"]})
        with st.chat_message("assistant"):
            st.markdown(result["answer"])
        st.session_state.chat_history.append((user_query, result["answer"]))