from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.documents import Document
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.llms import LlamaCpp
from langchain.chains import RetrievalQA

# 1. Load the local embedding model
embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# 2. Load and split the PDF
pdf_path = "my-notes.pdf"
loader = PyPDFLoader(pdf_path)
pages = loader.load()

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

# 3. Create the FAISS vector store
vectorstore = FAISS.from_documents(documents, embedding)

# 4. Load local LLaMA model
llm = LlamaCpp(
    model_path="models/llama-2-7b.Q4_K_M.gguf",  # 🔁 Change to your downloaded model path
    n_ctx=2048,
    temperature=0.5,
    top_p=0.95,
    verbose=True
)

# 5. Create RetrievalQA chain (RAG)
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectorstore.as_retriever(search_kwargs={"k": 2}),
    return_source_documents=False  # Set True if you want to see source chunks
)

# 6. Chat loop
print("\n🦙 PDF Chatbot using LLaMA is ready! Ask questions (type 'exit' to quit)\n")

while True:
    query = input("🧠 You: ")
    if query.lower() in ["exit", "quit"]:
        print("👋 Goodbye!")
        break

    response = qa_chain.run(query)
    print(f"🤖 Answer: {response}\n")