import os
import sys

import openai
from langchain.chains import ConversationalRetrievalChain, RetrievalQA
from langchain_community.chat_models import ChatOpenAI
from langchain_community.document_loaders import DirectoryLoader
from langchain_community.document_loaders import TextLoader
from langchain_openai import OpenAIEmbeddings
from langchain.indexes import VectorstoreIndexCreator
from langchain.indexes.vectorstore import VectorStoreIndexWrapper
from langchain_community.llms import OpenAI
from langchain_community.vectorstores import Chroma

import gradio as gr

  
# Enable to save to disk & reuse the model (for repeated queries on the same data)
PERSIST = False

chat_history=[];

query = None
if len(sys.argv) > 1:
  query = sys.argv[1]

if PERSIST and os.path.exists("persist"):
  print("Reusing index...\n")
  vectorstore = Chroma(persist_directory="persist", embedding_function=OpenAIEmbeddings())
  index = VectorStoreIndexWrapper(vectorstore=vectorstore)
else:
  #loader = TextLoader("data/data.txt") # Use this line if you only need data.txt
  loader = DirectoryLoader("data")
  if PERSIST:
    index = VectorstoreIndexCreator(vectorstore_kwargs={"persist_directory":"persist"}).from_loaders([loader])
  else:
    index = VectorstoreIndexCreator().from_loaders([loader])

chain = ConversationalRetrievalChain.from_llm(
  llm=ChatOpenAI(model="gpt-3.5-turbo-1106"),
  retriever=index.vectorstore.as_retriever(search_kwargs={"k": 1}),
)

  
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")


def answer_query(user_query):
    # Define your hidden context
    hidden_context = "o que acha a Iniciativa Liberal sobre: "
    # Prepend hidden context to the user's query
    full_query = hidden_context + user_query + "."
    # Now, use full_query instead of user_query to interact with the model
    result = chain({"question": full_query, "chat_history": chat_history})
    answer = result['answer']
    return answer

iface = gr.Interface(fn=answer_query,
                     inputs=gr.Textbox(label="Digite sua pergunta aqui", placeholder="Escreva aqui..."),
                     outputs=gr.Textbox(label="Resposta"),
                     title="Faz perguntas à Iniciativa Liberal",
                     description="Estas respostas são geradas com base no programa eleitoral da Iniciativa Liberal.")

iface.launch(share=True)