import gradio as gr


import os
hftoken = os.environ["hftoken"]

from langchain_huggingface import HuggingFaceEndpoint

repo_id = "mistralai/Mistral-7B-Instruct-v0.3"
llm = HuggingFaceEndpoint(repo_id = repo_id, max_new_tokens = 128, temperature = 0.7, huggingfacehub_api_token = hftoken)

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_template("tell me a joke about {topic}")
chain = prompt | llm | StrOutputParser()

# from langchain.document_loaders.csv_loader import CSVLoader
from langchain_community.document_loaders.csv_loader import CSVLoader


loader = CSVLoader(file_path='aiotsmartlabs_faq.csv', source_column = 'prompt')
data = loader.load()

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_huggingface.embeddings import HuggingFaceEndpointEmbeddings

# CHECK MTEB LEADERBOARD & FIND BEST EMBEDDING MODEL
model = "BAAI/bge-m3"
embeddings = HuggingFaceEndpointEmbeddings(model = model)

vectorstore = Chroma.from_documents(documents = data, embedding = embeddings)
retriever = vectorstore.as_retriever()

# from langchain.prompts import PromptTemplate

from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_template("""Given the following history, context and a question, generate an answer based on the context only.

In the answer try to provide as much text as possible from "response" section in the source document context without making much changes.
If somebody asks "Who are you?" or a similar phrase, state "I am Rishi's assistant built using a Large Language Model!"
If the answer is not found in the context, kindly state "I don't know. Please ask Rishi on Discord. Discord Invite Link: https://discord.gg/6ezpZGeCcM. Or email at rishi@aiotsmartlabs.com" Don't try to make up an answer.

CONTEXT: {context}

HISTORY: {history}

QUESTION: {question}""")

from langchain_core.runnables import RunnablePassthrough

# Define the chat response function
def chatresponse(message, history):
    history_langchain_format = []
    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))
    history_langchain_format.append(HumanMessage(content=message))
    gpt_response = llm(history_langchain_format)

    rag_chain = (
    {"context": retriever, "history": history_langchain_format, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
    )

    
    output = rag_chain.invoke(message)
    response = output.split('ANSWER: ')[-1].strip()
    return response

# Launch the Gradio chat interface
gr.ChatInterface(chatresponse).launch()

# import gradio as gr

# def chatresponse(message, history):
#     return history

# # Launch the Gradio chat interface
# gr.ChatInterface(chatresponse).launch()

# import gradio as gr
# from huggingface_hub import InferenceClient

# """
# For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
# """
# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")


# def respond(
#     message,
#     history: list[tuple[str, str]],
#     system_message,
#     max_tokens,
#     temperature,
#     top_p,
# ):
#     messages = [{"role": "system", "content": system_message}]

#     for val in history:
#         if val[0]:
#             messages.append({"role": "user", "content": val[0]})
#         if val[1]:
#             messages.append({"role": "assistant", "content": val[1]})

#     messages.append({"role": "user", "content": message})

#     response = ""

#     for message in client.chat_completion(
#         messages,
#         max_tokens=max_tokens,
#         stream=True,
#         temperature=temperature,
#         top_p=top_p,
#     ):
#         token = message.choices[0].delta.content

#         response += token
#         yield response

# """
# For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
# """
# demo = gr.ChatInterface(
#     respond,
#     additional_inputs=[
#         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
#         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
#         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
#         gr.Slider(
#             minimum=0.1,
#             maximum=1.0,
#             value=0.95,
#             step=0.05,
#             label="Top-p (nucleus sampling)",
#         ),
#     ],
# )


# if __name__ == "__main__":
#     demo.launch()