import gradio as gr
import fitz  # PyMuPDF
from langchain_community.llms import Ollama 
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate
from langchain.docstore.document import Document

def extract_text_from_pdf(file):
    doc = fitz.open(file.name)  # Fix lỗi: dùng tên file thay vì file.read()
    text = ""
    for page in doc:
        text += page.get_text()
    return text

# Trạng thái lưu lịch sử chat (nếu không dùng PDF)
chat_history = []

def answer_question(pdf_file, question, history):
    if not question:
        return "Vui lòng nhập câu hỏi.", history

    llm = Ollama(model="llama3")

    if pdf_file:  # Nếu có PDF được tải lên
        content = extract_text_from_pdf(pdf_file)
        prompt_template = PromptTemplate.from_template(
            "Trả lời câu hỏi sau dựa trên tài liệu:\n\n{context}\n\nCâu hỏi: {question}"
        )
        chain = load_qa_chain(llm, chain_type="stuff", prompt=prompt_template)
        docs = [Document(page_content=content)]
        result = chain.run(input_documents=docs, question=question)
        return result, history
    else:
        # Chat tự do (không có PDF)
        context = "\n".join([f"User: {q}\nAssistant: {a}" for q, a in history[-5:]])  # nhớ 5 lượt gần nhất
        prompt = f"{context}\nUser: {question}\nAssistant:"

        answer = llm(prompt)
        history.append((question, answer.strip()))
        return answer.strip(), history

# Giao diện Gradio nâng cấp
chatbot = gr.Chatbot()
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 PDF Q&A và Chat tự do với Ollama")
    with gr.Row():
        pdf_file = gr.File(label="📄 Tải PDF (tùy chọn)")
    with gr.Row():
        question = gr.Textbox(label="💬 Nhập câu hỏi", placeholder="Hỏi gì cũng được...")
    with gr.Row():
        submit_btn = gr.Button("Gửi câu hỏi")
    with gr.Row():
        chatbot_output = chatbot

    state = gr.State([])  # Lưu lịch sử hội thoại

    def respond(pdf, q, hist):
        answer, updated_history = answer_question(pdf, q, hist)
        return updated_history, updated_history

    submit_btn.click(respond, inputs=[pdf_file, question, state], outputs=[chatbot_output, state])

demo.launch()