import logging
import os

from llama_index.indices.document_summary import DocumentSummaryIndexEmbeddingRetriever
from llama_index.llms import OpenAI
from llama_index.query_engine import RetrieverQueryEngine

logging.basicConfig(format='%(asctime)s %(levelname)s:%(message)s', level=os.environ.get("LOGLEVEL", "DEBUG"))
import gradio as gr
from llama_index import VectorStoreIndex, StorageContext, download_loader, load_index_from_storage, ServiceContext, \
    get_response_synthesizer

cache = {}
chatgpt = OpenAI(temperature=0, model="gpt-3.5-turbo")
service_context = ServiceContext.from_defaults(llm=chatgpt, chunk_size=1024)

def loadData():
    index_root = "./summary_index"
    directory_names = os.listdir(index_root)
    for directory in directory_names:
        if os.path.isdir(f"{index_root}/{directory}"):
            print("Loading from existing index " + directory)
            storage_context = StorageContext.from_defaults(persist_dir=f"{index_root}/{directory}")
            index = load_index_from_storage(storage_context)
            retriever = DocumentSummaryIndexEmbeddingRetriever(
                index,
                # choice_select_prompt=choice_select_prompt,
                # choice_batch_size=choice_batch_size,
                # format_node_batch_fn=format_node_batch_fn,
                # parse_choice_select_answer_fn=parse_choice_select_answer_fn,
                service_context=service_context
            )
            # configure response synthesizer
            response_synthesizer = get_response_synthesizer(service_context=service_context)

            # assemble query engine
            query_engine = RetrieverQueryEngine(retriever=retriever, response_synthesizer=response_synthesizer)
            cache[directory] = query_engine


def chatbot(indexName, input_text):
    """
    Chatbot function that takes in a prompt and returns a response
    """
    response = cache[indexName].query(input_text)
    return response.response


def main():
    loadData()
    iface = gr.Interface(fn=chatbot, inputs=[
        gr.Dropdown(cache.keys(),
                    type="value", value="sos", label="Select Channel"),
        gr.Textbox(lines=7, label="Ask any question", placeholder='What are the key topics?')], outputs="text",
                         title="NLP Demo for Slack Data")
    if 'LOGIN_PASS' in os.environ:
        iface.launch(auth=('axiamatic', os.environ['LOGIN_PASS']),
                     auth_message='For access, please check my Slack profile or contact me in Slack.',
                     share=False)
    else:
        iface.launch(share=False)

main()