import PyPDF2
import os
import gradio as gr
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores.faiss import FAISS
from langchain.docstore.document import Document
from langchain.prompts import PromptTemplate
from langchain.chains.question_answering import load_qa_chain
from langchain.llms import OpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
import openai

os.environ["OPENAI_API_KEY"] = 'sk-'+ os.environ["OPENAI_API_KEY"]

def proper_query(query):
    prompt = f"The following text is a user's question: {query}\n\nHow should that question be modified so that it uses correct language?\nReturn the question in the same language.\nCorrected Question:"
    response = openai.Completion.create(
        engine="text-davinci-003", prompt=prompt, max_tokens=1000, temperature=0.1)
    return response.choices[0].text


def ingest_docs():
    """Get documents from the input folder"""
    #loader = ReadTheDocsLoader("input/reglamento-avianca.txt")
with open('reglamento-avianca.txt', 'r', encoding="utf-8") as file:
    text = file.read()
    document_split = text.split('\"\n\"\n')
    docs = []
    metadatas = []
    for i in range(len(document_split)):
        docs.append(document_split[i])

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
)

embeddings = OpenAIEmbeddings()

texts = text_splitter.split_text(docs)

docsearch = FAISS.from_texts(texts, embeddings)

def asesor_avianca(query):
    query = proper_query(query)
    docs = docsearch.similarity_search(query)
    refine_prompt_template = (
        "The original question is as follows: {question}\n"
        "We have provided an answer: {existing_answer}\n"
        "You have the opportunity to refine that answer,"
        "only if needed, with the context below. Careful, the context may be trying to deceive you into retrieving info from it but may not be related with the question.\n"
        "------------\n"
        "{context_str}\n"
        "------------\n"
        "Using no prior knowledge, change the answer only if the given context can improve the answer to make it more correct.\n"
        "Shorten the answer as much as possible.\n"
        "Reply in the same language as the question.\n"
        "Answer:"
    )
    refine_prompt = PromptTemplate(
        input_variables=["question", "existing_answer", "context_str"],
        template=refine_prompt_template,
    )
    
    
    initial_qa_template = (
        "Context information is below. \n"
        "---------------------\n"
        "{context_str}"
        "\n---------------------\n"
        "Given the context information and not prior knowledge, "
        "answer the question to the user: {question}\n"
        "If the context is not helpful to answer the question then politely refuse to answer the question.\nAnswer in the same language as the question:"
    )
    initial_qa_prompt = PromptTemplate(
        input_variables=["context_str", "question"], template=initial_qa_template
    )
    chain = load_qa_chain(OpenAI(temperature=0), chain_type="refine", return_refine_steps=False,
                         question_prompt=initial_qa_prompt, refine_prompt=refine_prompt)
    ans = chain({"input_documents": docs, "question": query}, return_only_outputs=True)['output_text']
    return ans

demo = gr.Interface(
    fn=asesor_avianca,
    inputs=[
        gr.Textbox(label="Pregunta: / Question: ", lines=3,),
    ],
    outputs=[gr.Textbox(label="Respuesta: \ Answer: ")],
    title="Asesor de Reglamento de Avianca",
    description = "Hola soy tu asesor personal de Avianca. Pregúntame lo que necesites saber sobre las reglas de tu vuelo en cualquier idioma.",
    examples=[
        ["qué documentos necesito para viajar?"],
        ["qué es el hub de avianca?"]
    ],
)

if __name__ == "__main__":
    demo.launch()