import streamlit as st
from dotenv import load_dotenv
from PyPDF2 import PdfReader
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.embeddings import HuggingFaceInstructEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.chat_models import ChatOpenAI
from langchain.llms import HuggingFaceHub
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
import os


def get_pdf_text(pdf_docs):
    text = ""
    for pdf in pdf_docs:
        pdf_reader = PdfReader(pdf)
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text

def get_text_chunks(text):
    text_splitter = CharacterTextSplitter(
        separator="\n",
        chunk_size=500, # the character length of the chunck
        chunk_overlap=100, # the character length of the overlap between chuncks
        length_function=len # the length function - in this case, character length (aka the python len() fn.)
    )
    chunks = text_splitter.split_text(text)
    return chunks

def get_vectorstore(text_chunks):
    model_name = "hkunlp/instructor-xl"
    hf = HuggingFaceInstructEmbeddings(model_name=model_name)
    vectorstore = FAISS.from_texts(texts=text_chunks, embedding=hf)
    return vectorstore

def get_conversation_chain(vectorstore):
    llm = HuggingFaceHub(repo_id="mistralai/Mistral-7B-Instruct-v0.2",model_kwargs={"Temperature": 0.5, "MaxTokens": 1024})
    retriever=vectorstore.as_retriever()
    prompt = hub.pull("rlm/rag-prompt")
    # Chain
    rag_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
    )
    response = rag_chain.invoke("A partir de documents PDF, concernant la transition écologique en France, proposer un plan de transition en fonction de la marque").split("\nAnswer:")[-1]
    return response

def rag_pdf():
    load_dotenv()
    st.header("SUGGESTION POUR PLAN DE TRANSITION DE MARQUE RESPECTANT LES NORMES DE LA RSE")

    if "conversation" not in st.session_state:
        st.session_state.conversation = None
    

    with st.sidebar:
        st.subheader("INFOS SUR LA MARQUE")
        pdf_docs = st.file_uploader("Upload les documents concerant la marque et clique sur process", type="pdf",accept_multiple_files=True)
        if st.button("Process"):
            with st.spinner("Processing..."):#loading bar to enhance user experience
                #get pdf text in raw format
                raw_text = get_pdf_text(pdf_docs)

                #get text chunks
                text_chunks = get_text_chunks(raw_text)

                #create vectorstore
                vectorstore = get_vectorstore(text_chunks)

                #create conversation chain
                st.session_state.conversation = get_conversation_chain(vectorstore)
    
    st.write(st.session_state.conversation)