import os
import pprint
import codecs
import chardet
import gradio
from langchain.llms import HuggingFacePipeline
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain import OpenAI, ConversationChain, LLMChain, PromptTemplate
from langchain.chains.conversation.memory import ConversationalBufferWindowMemory
from EdgeGPT import Chatbot


class VehicleManualChatbot:
    def __init__(self, db_paths, vehicle_options, embeddings):
        self.db_paths = db_paths
        self.vehicle_options = vehicle_options
        self.embeddings = embeddings
        self.chatgpt_chain = self._create_chatgpt_chain()

    def _create_chatgpt_chain(self, temperature=0.5):
        template = """
        {history}
        Human: {human_input}
        Assistant:"""

        prompt = PromptTemplate(
            input_variables=["history", "human_input"],
            template=template
        )

        return LLMChain(
            llm=OpenAI(temperature=temperature),
            prompt=prompt,
            verbose=True,
            memory=ConversationalBufferWindowMemory(k=2),
        )

    def _get_prompt(self, question, vehicle, k=4):
        prompt = f"""
        I need information from my {vehicle} manual.
        I will provide an excerpt from the manual. Use the excerpt and nothing else to answer the question.
        You must refer to the excerpt as "{vehicle} Manual" in your response. Here is the excerpt:
        """

        index = FAISS.load_local(
            folder_path=self.db_paths[vehicle], embeddings=self.embeddings)
        similar_docs = index.similarity_search(query=question, k=k)
        context = []
        for d in similar_docs:
            content = d.page_content
            context.append(content)
        user_input = prompt + '\n[EXCERPT]' + '\n' + \
            '\n'.join(context[:k]) + '\n' + 'Question:\n' + question
        return user_input

    def _ask_question(self, question, vehicle, k=2):
        index = FAISS.load_local(
            folder_path=self.db_paths[vehicle], embeddings=self.embeddings)

        prompt = self._get_prompt(question=question, vehicle=vehicle,
                                  k=k)
        response = self.chatgpt_chain.predict(human_input=prompt)
        return response

    def chat(self, question, vehicle, k=2, temperature=0.5):
        self.chatgpt_chain = self._create_chatgpt_chain(
            temperature=temperature)
        response = self._ask_question(question=question, vehicle=vehicle,
                                      k=k)
        return response


db_paths = {
    "S-Class": "data/s-class-manual",
    "EQS": "data/eqs-manual",
    "EQE": "data/eqe-manual",
}

embeddings = HuggingFaceEmbeddings()

vehicle_options = ["S-Class", "EQS", "EQE"]
chatbot = VehicleManualChatbot(db_paths=db_paths,
                               vehicle_options=vehicle_options,
                               embeddings=embeddings)


def start_ui():
    chatbot_interface = gradio.Interface(
        fn=chatbot.chat,
        inputs=["text",
                gradio.inputs.Dropdown(
                    vehicle_options, label="Select Vehicle Model"),
                gradio.inputs.Slider(minimum=1, maximum=10, step=1, label="k")
                ],
        outputs="text",
        title="Owner's Manual",
        description="Ask your vehicle manual and get a response.",
        examples=[["What are the different features of the dashboard console?", "S-Class", 2],
                  ["What is flacon?", "S-Class", 3],
                  ["What is hyperscreen?", "EQS", 2],
                  ["Where can I find my vin?", "EQS", 3]
                  ]
    )

    chatbot_interface.launch()