import gradio as gr
import os
from dotenv import load_dotenv
from langchain.chains.question_answering import load_qa_chain
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

load_dotenv()
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

llm = ChatOpenAI(model_name="gpt-4", temperature=0)
embedding = OpenAIEmbeddings()

db = None  # Global DB

def upload_pdf(file):
    global db
    loader = PyPDFLoader(file.name)
    documents = loader.load()
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs = splitter.split_documents(documents)
    db = FAISS.from_documents(docs, embedding)
    return "✅ PDF processed. Ask your questions below."

def chat_with_pdf(message, history):
    global db
    if not db:
        return history + [[message, "❌ Please upload a PDF first."]]
    
    retriever = db.as_retriever()
    chain = load_qa_chain(llm, chain_type="stuff")
    docs = retriever.get_relevant_documents(message)
    response = chain.run(input_documents=docs, question=message)
    history.append([message, response])
    return history

with gr.Blocks(title="📄 Interactive PDF Chatbot") as demo:
    gr.Markdown("## Talk to a PDF")

    with gr.Row():
        pdf = gr.File(label="Upload PDF", file_types=[".pdf"])
        upload_btn = gr.Button("Process PDF")
    status = gr.Textbox(label="Status", interactive=False)

    chatbot = gr.Chatbot(label="Chat History", height=400)
    msg = gr.Textbox(label="Ask a question", placeholder="Type your question and press Enter...", lines=2)
    send_btn = gr.Button("Send")

    upload_btn.click(upload_pdf, inputs=pdf, outputs=status)
    msg.submit(chat_with_pdf, inputs=[msg, chatbot], outputs=chatbot)
    send_btn.click(chat_with_pdf, inputs=[msg, chatbot], outputs=chatbot)
    send_btn.click(lambda: "", None, msg)  # Clear input after send
    msg.submit(lambda: "", None, msg)      # Clear input after enter

demo.launch(share=True)