import gradio as gr

from langchain.document_loaders import OnlinePDFLoader

from langchain.text_splitter import CharacterTextSplitter

from langchain.llms import HuggingFaceHub

from langchain.embeddings import HuggingFaceHubEmbeddings

from langchain.vectorstores import Chroma

from langchain.chains import RetrievalQA

import os
import tempfile
import openai
import json
import re
from langchain.docstore.document import Document
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI
from langchain.document_loaders import PyPDFLoader
from langchain.indexes import VectorstoreIndexCreator
import tempfile
from langchain.document_loaders import DirectoryLoader

os.environ["OPENAI_API_KEY"] = os.environ['OpenApi_Key']
query1=" "
limit = 0
def loading_pdf():
    return "Loading..."

def remove_last_slash_and_characters(input_string):
    last_slash_index = input_string.rfind('/')
    if last_slash_index != -1:
        result_string = input_string[:last_slash_index]
    else:
        result_string = input_string

    return result_string

def pdf_changes(pdf_doc, prompt):
    
    loader = DirectoryLoader(remove_last_slash_and_characters(pdf_doc.name))
    #loader = OnlinePDFLoader(pdf_doc.name)
    #loader = PyPDFLoader(pdf_doc.name)
    documents = loader.load()
    print(len(documents))

    #name_filter = "**/*.md"
    separator = "\n"
    chunk_size_limit = 1000
    max_chunk_overlap = 50

    text_splitter = CharacterTextSplitter(separator=separator, chunk_size=chunk_size_limit, chunk_overlap=max_chunk_overlap)
    split_docs = text_splitter.split_documents(documents)

    embeddings = OpenAIEmbeddings()
    vector_store1 = FAISS.from_documents(split_docs, embeddings)

    from langchain.prompts import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
      )


    system_template="""You are a helpful chatbot used by the user to chat with pdf documents. Only answer the questions by using information provided in the context provided to you. If there is no relavant context, tell 'Hmm, I'm not sure'."""+prompt+"""
    ----------------
    {summaries}"""

    messages = [
        SystemMessagePromptTemplate.from_template(system_template),
        HumanMessagePromptTemplate.from_template("{question}")
    ]
    prompt2 = ChatPromptTemplate.from_messages(messages)

    from langchain.chat_models import ChatOpenAI
    from langchain.chains import RetrievalQAWithSourcesChain
    global query1
    chain_type_kwargs = {"prompt": prompt2}
    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0, max_tokens=512)
    global chain
    chain = RetrievalQAWithSourcesChain.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store1.as_retriever(search_kwargs={'k': 2}),
        return_source_documents=True,
        chain_type_kwargs=chain_type_kwargs
    )
    return "Ready"

def add_text(history, text):
    history = history + [(text, None)]
    return history, ""

def bot(history):
    response = infer(history[-1][0])
    history[-1][1] = response
    return history

def infer(question):
    global query1
    global limit
    openai.api_key = os.environ['OpenApi_Key']
    prompt_text = question
    if prompt_text:
        query1 = query1 + "\nUser: " + prompt_text + "\nBot: "
        result = chain(query1)
        query1 = query1 + result['answer']
        query1 = openai.ChatCompletion.create(
                            model="gpt-3.5-turbo",
                            messages=[
                            {"role": "system", "content": "You are provided with chat history and latset conversation between user and bot. Summarise the history and latest conversationin minimum most tokens possible. Do not include greetings in the summary like hi, hello, etc."},
                            {"role": "user", "content": query1},
                                ]
                                ) ["choices"][0]["message"]["content"].replace("'", "")
        return result['answer']

css="""
#col-container { margin-left: auto; margin-right: auto;}
"""

title = """
<div style="text-align: center; max-width: 700px;">
    <h1 style="color: #4545FF;">Chat with PDF</h1>
    <p style="text-align: center; color: #4545FF;">Upload a .PDF from your computer, click the "Load PDF" button, <br />
    when everything is ready, you can start asking questions about the pdf ;)</p>
</div>
"""


with gr.Blocks(css=css,theme = gr.themes.Soft()) as demo:
    with gr.Column(elem_id="col-container"):
        #gr.HTML(title)
        with gr.Row():
          with gr.Column(scale=1):
            #gr.File(file_count="multiple")
              pdf_doc = gr.File(label="Load a pdf", file_count="single")
              prompt = gr.Textbox(label="Behaviour Prompt (optional)", placeholder="Reply to all questions as a rap / Reply to all questions in Hindi etc. ")
              #repo_id = gr.Dropdown(label="LLM", choices=["google/flan-ul2", "OpenAssistant/oasst-sft-1-pythia-12b", "bigscience/bloomz"], value="google/flan-ul2")
              with gr.Row():
                  langchain_status = gr.Textbox(label="Status", placeholder="Waiting for PDF", interactive=False,show_label=False)
                  load_pdf = gr.Button("Load pdf")
          with gr.Column(scale=2):
            chatbot = gr.Chatbot([], elem_id="chatbot",show_label=False,show_share_button=False).style(height=750)
            with gr.Row():
              question = gr.Textbox(label="Question", placeholder="Type your question and hit Enter ",scale=6,show_label=False)
              submit_btn = gr.Button("Send",scale=1)
    load_pdf.click(pdf_changes, inputs=[pdf_doc,prompt], outputs=[langchain_status], queue=False)
    question.submit(add_text, [chatbot, question], [chatbot, question]).then(
        bot, chatbot, chatbot
    )
    submit_btn.click(add_text, [chatbot, question], [chatbot, question]).then(
        bot, chatbot, chatbot
    )

demo.launch(debug=True)