import os
import pickle
import langchain

import faiss
from langchain import HuggingFaceHub
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import DirectoryLoader, TextLoader, UnstructuredHTMLLoader
from langchain.embeddings import OpenAIEmbeddings, HuggingFaceHubEmbeddings
from langchain.memory import ConversationBufferWindowMemory
from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
)
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores.faiss import FAISS
from langchain.cache import InMemoryCache

langchain.llm_cache = InMemoryCache()

global model_name

models = ["GPT-3.5", "Flan UL2", "GPT-4", "Flan T5"]

pickle_file = "_vs.pkl"
index_file = "_vs.index"
models_folder = "models/"

llm = ChatOpenAI(model_name="gpt-4", temperature=0.1)

embeddings = OpenAIEmbeddings(model='text-embedding-ada-002')

chat_history = []

memory = ConversationBufferWindowMemory(memory_key="chat_history", k=10)

vectorstore_index = None

system_template = """You are Coursera QA Bot. Have a conversation with a human, answering the following questions as best you can.
You are a teaching assistant for a Coursera Course: The 3D Printing Evolution and can answer any question about that using vectorstore or context.
Use the following pieces of context to answer the users question. 
----------------
{context}"""

messages = [
    SystemMessagePromptTemplate.from_template(system_template),
    HumanMessagePromptTemplate.from_template("{question}"),
]
CHAT_PROMPT = ChatPromptTemplate.from_messages(messages)


def set_model_and_embeddings(model):
    global chat_history
    set_model(model)
    # set_embeddings(model)
    chat_history = []


def set_model(model):
    global llm
    print("Setting model to " + str(model))
    if model == "GPT-3.5":
        print("Loading GPT-3.5")
        llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.1)
    elif model == "GPT-4":
        print("Loading GPT-4")
        llm = ChatOpenAI(model_name="gpt-4", temperature=0.1)
    elif model == "Flan UL2":
        print("Loading Flan-UL2")
        llm = HuggingFaceHub(repo_id="google/flan-ul2", model_kwargs={"temperature": 0.1, "max_new_tokens":500})
    elif model == "Flan T5":
        print("Loading Flan T5")
        llm = HuggingFaceHub(repo_id="google/flan-t5-base", model_kwargs={"temperature": 0.1})
    else:
        print("Loading GPT-3.5 from else")
        llm = ChatOpenAI(model_name="text-davinci-002", temperature=0.1)


def set_embeddings(model):
    global embeddings
    if model == "GPT-3.5" or model == "GPT-4":
        print("Loading OpenAI embeddings")
        embeddings = OpenAIEmbeddings(model='text-embedding-ada-002')
    elif model == "Flan UL2" or model == "Flan T5":
        print("Loading Hugging Face embeddings")
        embeddings = HuggingFaceHubEmbeddings(repo_id="sentence-transformers/all-MiniLM-L6-v2")


def get_search_index(model):
    global vectorstore_index
    if os.path.isfile(get_file_path(model, pickle_file)) and os.path.isfile(
            get_file_path(model, index_file)) and os.path.getsize(get_file_path(model, pickle_file)) > 0:
        # Load index from pickle file
        with open(get_file_path(model, pickle_file), "rb") as f:
            search_index = pickle.load(f)
            print("Loaded index")
    else:
        search_index = create_index(model)
        print("Created index")

    vectorstore_index = search_index
    return search_index


def create_index(model):
    source_chunks = create_chunk_documents()
    search_index = search_index_from_docs(source_chunks)
    faiss.write_index(search_index.index, get_file_path(model, index_file))
    # Save index to pickle file
    with open(get_file_path(model, pickle_file), "wb") as f:
        pickle.dump(search_index, f)
    return search_index


def get_file_path(model, file):
    # If model is GPT3.5 or GPT4 return models_folder + openai + file else return models_folder + hf + file
    if model == "GPT-3.5" or model == "GPT-4":
        return models_folder + "openai" + file
    else:
        return models_folder + "hf" + file


def search_index_from_docs(source_chunks):
    # print("source chunks: " + str(len(source_chunks)))
    # print("embeddings: " + str(embeddings))

    search_index = FAISS.from_documents(source_chunks, embeddings)
    return search_index


def get_html_files():
    loader = DirectoryLoader('docs', glob="**/*.html", loader_cls=UnstructuredHTMLLoader, recursive=True)
    document_list = loader.load()
    return document_list


def fetch_data_for_embeddings():
    document_list = get_text_files()
    document_list.extend(get_html_files())
    print("document list: " + str(len(document_list)))
    return document_list


def get_text_files():
    loader = DirectoryLoader('docs', glob="**/*.txt", loader_cls=TextLoader, recursive=True)
    document_list = loader.load()
    return document_list


def create_chunk_documents():
    sources = fetch_data_for_embeddings()

    splitter = CharacterTextSplitter(separator=" ", chunk_size=800, chunk_overlap=0)

    source_chunks = splitter.split_documents(sources)

    print("chunks: " + str(len(source_chunks)))

    return source_chunks


def get_qa_chain(vectorstore_index):
    global llm, model_name
    print(llm)

    # embeddings_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.76)
    # compression_retriever = ContextualCompressionRetriever(base_compressor=embeddings_filter, base_retriever=gpt_3_5_index.as_retriever())
    retriever = vectorstore_index.as_retriever(search_type="similarity_score_threshold",
                                               search_kwargs={"score_threshold": .5})

    chain = ConversationalRetrievalChain.from_llm(llm, retriever, return_source_documents=True,
                                                  verbose=True, get_chat_history=get_chat_history,
                                                  combine_docs_chain_kwargs={"prompt": CHAT_PROMPT})
    return chain


def get_chat_history(inputs) -> str:
    res = []
    for human, ai in inputs:
        res.append(f"Human:{human}\nAI:{ai}")
    return "\n".join(res)


def generate_answer(question) -> str:
    global chat_history, vectorstore_index
    chain = get_qa_chain(vectorstore_index)

    result = chain(
        {"question": question, "chat_history": chat_history, "vectordbkwargs": {"search_distance": 0.6}})
    chat_history = [(question, result["answer"])]
    sources = []
    print(result)

    for document in result['source_documents']:
        source = document.metadata['source']
        sources.append(source.split('/')[-1].split('.')[0])
        print(sources)

    source = ',\n'.join(set(sources))
    return result['answer'] + '\nSOURCES: ' + source