import os
import streamlit as st
from PyPDF2 import PdfReader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains.question_answering import load_qa_chain
import random
from langchain import HuggingFaceHub
from langchain.callbacks import get_openai_callback

def main():
    # ---------------------------- created personal API -----------------------------
    os.environ["HUGGINGFACEHUB_API_TOKEN"] = "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"

    # ------------------ Designing Page ---------------
    st.set_page_config(page_title="Ask Your PDF")
    st.header("Ask your PDF :")

    pdf = st.file_uploader("Upload your File here", type="pdf")

    # Check Pdf
    if pdf is not None:
        pdf_reader = PdfReader(pdf)

        text = ""

        # Extract pages from pdf
        for page in pdf_reader.pages:
            text += page.extract_text()

        # split into chunks
        text_spliter = CharacterTextSplitter(
            separator="\n",
            chunk_size=1000,
            chunk_overlap=0,
            length_function=len
        )

        chunks = text_spliter.split_text(text)

        # create embeddings
        embedding = HuggingFaceEmbeddings()
        knowledge_base = FAISS.from_texts(chunks, embedding)

        user_questions = st.text_input("Ask a Question from PDF : ")
        if user_questions:

            greeting = ["hy", 'hello', 'hey', "hi"]
            greet_msg = ["Hello Dear!", 'Hey!', 'Hey Friend!']

            if user_questions in greeting:
                response = random.choice(greet_msg)
            elif user_questions == "by" or user_questions == "bye":
                response = "GoodBye Sir!, Have a Nice Day....."
            else:

                docs = knowledge_base.similarity_search(user_questions)

                chain = load_qa_chain(HuggingFaceHub(repo_id="google/flan-t5-xxl", model_kwargs={"temperature":0.1, "max_length":512}), chain_type="stuff")
                with get_openai_callback() as cb:

                    response = chain.run(input_documents=docs, question=user_questions)
                    print(cb)

            st.write(response)


if __name__ == "__main__":
    main()