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.\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 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 refuse to answer the question in the same language."
    )
    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, Hexagonito. Pregúntame lo que necesites saber sobre las reglas de tu vuelo. \nHi, I am Hexagonito, your Avianca rules personal assistant, ask me anything about the rules applying to your flight in any language.",
    examples=[
        ["qué documentos necesito para viajar?"],
        ["no llegó mi equipaje, qué puedo hacer?"]
    ],
)

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