import logging
import os
import json
from llama_index.indices.document_summary import DocumentSummaryIndexEmbeddingRetriever
from llama_index.indices.vector_store import VectorIndexRetriever
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", "INFO"))
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 load_mapping_from_json(filepath):
    if not os.path.exists(filepath):
        return {}
    with open(filepath, 'r') as file:
        return json.load(file)

userIdMapping = load_mapping_from_json('user_id_to_name_mapping.json')

def loadData():
    index_root = "./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)
            vector_retriever = VectorIndexRetriever(index=index, similarity_top_k=5)
            response_synthesizer = get_response_synthesizer(service_context=service_context,
                                                            response_mode="tree_summarize")
            query_engine = RetrieverQueryEngine(
                retriever=vector_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
    """
    query = "This data contains updates from multiple teams and sprints - "+ input_text + " Ignore headers in the content."
    response = cache[indexName].query(query)
    answer = response.response.replace('Based on the given context information', 'Based on the available information')
    for userId in userIdMapping:
        answer = answer.replace(userId, userIdMapping[userId])
    return answer


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=gr.Textbox(lines=11, label="Response"),
                         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()