import gradio
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.vectorstores import FAISS
from langchain import OpenAI, ConversationChain, LLMChain, PromptTemplate
from langchain.chains.conversation.memory import ConversationBufferMemory

def chat(question, vehicle, k=10, temperature=0.01):
    chatgpt_chain = create_chatgpt_chain(temperature=temperature)
    response = ask_question(question=question, vehicle=vehicle, k=k, embeddings=model_norm, chatgpt_chain=chatgpt_chain)
    return response

def create_chatgpt_chain(temperature):
    template = """
    {chat_history}
    Human: {question}
    AI:
    """
    prompt_template = PromptTemplate(input_variables=["chat_history", "question"], template=template)
    return LLMChain(llm=OpenAI(temperature=temperature,model_name="gpt-3.5-turbo"),prompt=prompt_template,verbose=True,memory=ConversationBufferMemory(memory_key="chat_history"))

def ask_question(question, vehicle, k, embeddings, chatgpt_chain):
    index = FAISS.load_local(folder_path=db_paths[vehicle], embeddings=embeddings)
    prompt = get_prompt(question=question, vehicle=vehicle, k=k)
    response = chatgpt_chain.run(question=prompt)
    return response

def get_prompt(question, vehicle, k):
    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=db_paths[vehicle], embeddings=model_norm)
    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'.join(context[:k]) + '\nQuestion: ' + question
    return user_input

db_paths = {
  "2023 AMG C-Coupe-Cab": "data/amg_c_coupe_cab",
  "2023 AMG C-Sedan": "data/amg_c_sedan",
  "2023 AMG E-Coupe-Cab": "data/amg_e_coupe_cab",
  "2023 AMG E-Sedan_wagon": "data/amg_e_sedan_wagon",
  "2023 AMG_EQE-Sedan": "data/amg_eqe_sedan",
  "2023 AMG_GLE-suv": "data/amg_gle_suv",
  "2023 AMG_GLS SUV": "data/amg_gls_suv",
  "2023 C-Cab": "data/c_cab",
  "2023 C-Coupe": "data/c_coupe",
  "2023 C-Sedan": "data/c_sedan",
  "2023 CLA": "data/cla",
  "2023 E-Cab": "data/e_cab",
  "2023 E-Coupe": "data/e_coupe",
  "2023 E-Sedan": "data/e_sedan",
  "2023 E-wagon": "data/e_wagon",
  "2023 eqb SUV": "data/eqb_suv",
  "2023 EQE-Sedan": "data/eqe_sedan",
  "2023 EQS_Sedan": "data/eqs_sedan",
  "2023 EQS SUV": "data/eqs_suv",
  "2023 GLA": "data/gla",
  "2023 GLB": "data/glb",
  "2023 GLC-Coupe": "data/glc_coupe",
  "2023 GLE-Coupe": "data/gle_coupe",
  "2023 GLE-suv": "data/gle_suv",
  "2023 GLS SUV": "data/gls_suv"
}

vehicle_options = [
    "2023 AMG C-Coupe-Cab",
    "2023 AMG C-Sedan",
    "2023 AMG E-Coupe-Cab",
    "2023 AMG E-Sedan_wagon",
    "2023 AMG_EQE-Sedan",
    "2023 AMG_GLE-suv",
    "2023 AMG_GLS SUV",
    "2023 C-Cab",
    "2023 C-Coupe",
    "2023 C-Sedan",
    "2023 CLA",
    "2023 E-Cab",
    "2023 E-Coupe",
    "2023 E-Sedan",
    "2023 E-wagon",
    "2023 eqb SUV",
    "2023 EQE-Sedan",
    "2023 EQS SUV",
    "2023 EQS_Sedan",
    "2023 GLA",
    "2023 GLB",
    "2023 GLC-Coupe",
    "2023 GLE-Coupe",
    "2023 GLE-suv",
    "2023 GLS SUV",
]

model_name = "BAAI/bge-large-en"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
model_norm = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)

def start_ui():
    chatbot_interface = gradio.Interface(
        fn=chat,
        inputs=["text",
                gradio.inputs.Dropdown(vehicle_options, label="Select Mercedes-Benz Owner's Manual")
                #gradio.inputs.Slider(minimum=1, maximum=10, step=1, label="k")
        ],
        outputs="text",
        title="Mercedes-Benz Owner's Manual",
        description="Ask a question and get answers from Mercedes-Benz Owner's Manual.<u>Disclaimer:</u> THIS IS NOT OFFICIAL AND MAY NOT BE AVAILABLE ALL THE TIME. ALWAYS LOOK AT THE OFFICIAL DOCUMENTATION at https://www.mbusa.com/en/owners/manuals",
        examples=[["What are the different features of the dashboard console?", "2023 S-Class", 10, 0.01],
                  ["What is flacon? Which page has that information? Show me all the exact content from that page", "2023 S-Class", 10, 0.01],
                  ["What is hyperscreen?", "2023 EQS", 10, 0.01],
                  ["Where can I find my vin?", "2023 EQS", 10, 0.01],
                  ["Does it take more than 30 minutes to charge? Which page has that information? Show me all the exact content from that page", "2023 EQE", 10, 0.01]],
        article = '<center><img src="https://visitor-badge.glitch.me/badge?page_id=kaushikdatta.owner-manual" alt="visitor badge"/></center>'
    )

    chatbot_interface.launch()