from langchain_ollama.llms import OllamaLLM
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain
from langchain_huggingface import HuggingFaceEmbeddings
import pandas as pd
import os
from langchain.schema import Document
def load_and_process_pdf():
    df = pd.read_excel(os.path.join(os.getcwd(), "chunk_metadata_template.xlsx"))
    splits = []
    for i, row in df.iterrows():
        # Create a Document object for each row, including page_content and metadata
        chunk_with_metadata = Document(
            page_content=row['page_content'],  # Content for the chunk
            metadata={
                'chunk_id': row['chunk_id'],  # Add chunk_id to the metadata
                'document_title': row['document_title'],  # Add document_title to the metadata
                'keywords': row['keywords'],  # Add keywords to the metadata
            }
        )
        # Append the Document object to the chunks list
        splits.append(chunk_with_metadata)
    return splits

def create_vectorstore(splits = load_and_process_pdf()):
    # model_name = "nomic-ai/nomic-embed-text-v1"
    embeddings = HuggingFaceEmbeddings()
    #     model_name=model_name)
    vectorstore = Chroma(
        collection_name="example_collection",
        embedding_function=embeddings,
        persist_directory="./chroma_langchain_db"  # Optional: Directory for local persistence
    )
    vectorstore.add_documents(splits)  # Add processed document splits to the vector store
    return vectorstore

def create_rag_chain():
    llm = OllamaLLM(model="gemma2:2b")
    
    system_prompt = (
        """
<|context|>
Bạn là một AI trợ lý chuyên cung cấp về sản phẩm bảo hiểm của công ty MB Ageas Life tại Việt Nam.
Hãy trả lời  chuyên nghiệp, chính xác, cung cấp thông tin bao quát trước, đưa ví dụ rồi mới đặt câu hỏi gợi mở nếu chưa rõ. Tất cả các thông tin cung cấp đều trong phạm vi MBAL. Những câu hỏi không có đủ thông tin thì mời khách hàng đăng ký để nhận tư vấn trên https://www.mbageas.life/
</s>
<|user|>
{query}
</s>
 <|assistant|>
"""
    )
    
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            ("human", "{input}"),
        ]
    )
    
    question_answer_chain = create_stuff_documents_chain(llm, prompt)
    return question_answer_chain

def get_response(rag_chain, vectorstore, question):
    retriever = vectorstore.as_retriever(
    search_kwargs={'k': 5}
)
    retrieval_chain = create_retrieval_chain(retriever, rag_chain)
    response = retrieval_chain.invoke({"input": question})
    return response["answer"]