import os
import gradio as gr
from langchain.chains import ConversationalRetrievalChain
from langchain.text_splitter import CharacterTextSplitter
from langchain.document_loaders import PyPDFLoader
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.embeddings.openai import OpenAIEmbeddings

vectordb = None

# 處理 PDF 文件並初始化向量數據庫
def load_and_process_pdf(pdf_file, api_key):
    try:
        loader = PyPDFLoader(pdf_file.name)
        documents = loader.load()
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
        docs = text_splitter.split_documents(documents)

        # 傳遞 API Key
        embeddings = OpenAIEmbeddings(openai_api_key=api_key)
        vectordb = Chroma.from_documents(docs, embedding=embeddings, persist_directory="./data")
        return vectordb
    except Exception as e:
        return f"無法處理文件，錯誤: {str(e)}"

# 查詢處理函數
def handle_query(api_key, user_message, pdf_vectordb, chat_history):
    if not pdf_vectordb or isinstance(pdf_vectordb, str):
        return chat_history, "請先上傳有效的 PDF 文件。"

    try:
        retriever = pdf_vectordb.as_retriever(search_kwargs={"k": 5})
        pdf_qa_chain = ConversationalRetrievalChain.from_llm(
            ChatOpenAI(temperature=0.7, model_name="gpt-4", openai_api_key=api_key),
            retriever=retriever
        )
        result = pdf_qa_chain({"question": user_message, "chat_history": chat_history})
        answer = result.get("answer", "抱歉，我無法提供答案。")
        chat_history.append((user_message, answer))
        return chat_history, answer
    except Exception as e:
        return chat_history, f"出現錯誤: {str(e)}"

# 當用戶上傳 PDF 時處理
def process_pdf_upload(api_key, pdf_file):
    global vectordb
    vectordb = load_and_process_pdf(pdf_file, api_key)
    if isinstance(vectordb, str):
        return vectordb
    return "食譜已成功載入！請開始點菜。"

# 主程式 - Gradio 介面
with gr.Blocks() as demo:
    gr.Markdown("## 📘 國際大廚教你做菜 ")

    with gr.Row():
        api_key = gr.Textbox(label="🔑 請輸入您的 OpenAI API Key", type="password")
        pdf_file = gr.File(label="📂 上傳 PDF 食譜", file_types=[".pdf"])

    chatbot = gr.Chatbot(label="💬 聊天區")
    state = gr.State([])

    with gr.Row():
        user_input = gr.Textbox(show_label=False, placeholder="請輸入問題...", lines=2)
        send_btn = gr.Button("提交問題")

    response = gr.Textbox(label="📢 回應", interactive=False, lines=4)

    pdf_file.change(process_pdf_upload, inputs=[api_key, pdf_file], outputs=response)

    def handle_user_input(api_key, user_message, chat_history):
        if not vectordb:
            return chat_history, "請先上傳 PDF 食譜。"
        if user_message.strip().lower() == "謝謝":
            return chat_history, "加油~~~"
        return handle_query(api_key, user_message, vectordb, chat_history)

    send_btn.click(handle_user_input, inputs=[api_key, user_input, state], outputs=[chatbot, response])
    user_input.submit(handle_user_input, inputs=[api_key, user_input, state], outputs=[chatbot, response])

demo.launch()