import os
import chainlit as cl
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_groq import ChatGroq  # ✅ Correct import
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA

# PDF to QA chain processor
def process_file(file_path):
    loader = PyPDFLoader(file_path)
    documents = loader.load()
    cl.run_sync(cl.Message("📄 PDF loaded successfully. Splitting text...").send())

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    texts = text_splitter.split_documents(documents)

    cl.run_sync(cl.Message("🔍 Creating embeddings and FAISS vector store...").send())
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vectorstore = FAISS.from_documents(texts, embeddings)

    cl.run_sync(cl.Message("⚙️ Building RetrievalQA chain...").send())

    prompt_template = """
    Use the following pieces of context to answer the question at the end.
    If you don't know the answer, just say you don't know — don't try to make up an answer.
    
    {context}
    
    Question: {question}
    Helpful Answer:
    """

    prompt = PromptTemplate(
        template=prompt_template,
        input_variables=["context", "question"]
    )

    # ✅ Use Groq Chat Model
    llm = ChatGroq(
        api_key=os.environ.get("GROQ_API_KEY"),
        model_name="llama3-8b-8192"
    )

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
        chain_type_kwargs={"prompt": prompt},
        return_source_documents=True
    )

    return qa_chain

# On Chat Start
@cl.on_chat_start
async def start():
    await cl.Message("👋 Welcome! Upload a PDF to begin.").send()

    await cl.AskFileMessage(
        content="📄 Upload a PDF file below (Max 20 MB):",
        accept=["application/pdf"],
        max_size_mb=20,
        timeout=180
    ).send()

# On User Message
@cl.on_message
async def handle_message(message: cl.Message):
    if files := await cl.user_session.get("files"):
        file = files[0]
        file_path = file.path

        qa_chain = await cl.make_async(process_file)(file_path)
        cl.user_session.set("qa_chain", qa_chain)
        cl.user_session.set("ready", True)

        await cl.Message(f"✅ File `{file.name}` processed. QA chain is ready. You can now ask your question!").send()
        cl.user_session.set("files", None)
        return

    if not cl.user_session.get("ready"):
        await cl.Message("⚠️ Please upload a PDF file first.").send()
        return

    qa_chain = cl.user_session.get("qa_chain")
    await cl.Message("💬 Thinking...").send()

    response = qa_chain(message.content)
    answer = response["result"]
    await cl.Message(content=f"🧠 Answer: {answer}").send()