from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains.question_answering import load_qa_chain
from langchain.llms import OpenAI
from langchain.document_loaders import BSHTMLLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
import os

import gradio as gr
import locale
locale.getpreferredencoding = lambda: "UTF-8"

print("LOGGING")

# Load the files
directory = './data/'
#bshtml_dir_loader = DirectoryLoader(directory, loader_cls=BSHTMLLoader,loader_kwargs={'features': 'html.parser'})
bshtml_dir_loader = DirectoryLoader(directory, loader_cls=lambda path: BSHTMLLoader(path, bs_kwargs={'features': 'html.parser'}))
data = bshtml_dir_loader.load()
#Split the document into chunks
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 1000,
    chunk_overlap  = 20,
    length_function = len,
)

documents = text_splitter.split_documents(data)

print("Got docs split")

# Create the embeddings
embeddings = OpenAIEmbeddings()
#Load the model
llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo")

# Create vectorstore to use as the index
vectordb = Chroma.from_documents(documents=documents, embedding=embeddings)

#expose this index in a retriever object
doc_retriever = vectordb.as_retriever()
print("Created retriever")
#create the QA chain
ted_lasso_qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=doc_retriever)
# Function to make inferences and provide answers
def make_inference(query):
    print("reached inference")
    return ted_lasso_qa.run(query)

if __name__ == "__main__":
    # make a gradio interface
    import gradio as gr

    gr.Interface(
        make_inference,
        [
            gr.inputs.Textbox(lines=2, label="Query"),
        ],
        gr.outputs.Textbox(label="Response"),
        title="Ask me about Ted Lasso 📺⚽",
        description="Ask me about Ted Lasso 📺⚽ is a tool that allows you to ask questions the tv series Ted Lasso",
    ).launch()