#Reference : https://medium.com/@tahreemrasul/building-a-chatbot-application-with-chainlit-and-langchain-3e86da0099a6
#Reference : https://platform.deepseek.com/api-docs/api/create-chat-completion 
from langchain.chains import LLMChain
from prompts import maths_assistant_prompt_template
from langchain.memory.buffer import ConversationBufferMemory
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import torch
import chainlit as cl
#from chainlit import Button  # Import Button component
# Load the model and tokenizer
model_name = "deepseek-ai/deepseek-math-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
model.generation_config = GenerationConfig.from_pretrained(model_name)
model.generation_config.pad_token_id = model.generation_config.eos_token_id


@cl.on_chat_start
async def start_llm():

    conversation_memory = ConversationBufferMemory(memory_key="chat_history",
                                                   max_len=50,
                                                   return_messages=True,
                                                   )
    llm_chain = LLMChain(llm=model, prompt=maths_assistant_prompt_template, memory=conversation_memory)
    cl.user_session.set("llm_chain", llm_chain)

    #Send initial message to the user
    #await cl.Message("What is your topic of interest?").send()
    # Send initial message with selectable buttons
    actions = [
        cl.Action(name="Probability", value="Probability", description="Select Quiz Topic!"),
        cl.Action(name="Linear Algebra", value="Linear Algebra", description="Select Quiz Topic!"),
        cl.Action(name="Accounts", value="Accounts", description="Select Quiz Topic!"),
        cl.Action(name="Calculus", value="Calculus", description="Select Quiz Topic!")
    ]
    await cl.Message(content="**Pick a Topic and Let the Quiz Adventure Begin!** 🎉📚", actions=actions).send()



@cl.on_message
async def query_llm(message: cl.Message):
    llm_chain = cl.user_session.get("llm_chain")
    #selected_topic = cl.user_session.get("selected_topic", "probability")  # Default to probability if not set
    print("Message being sent to the LLM is")
    print(message.content)
    #response = await llm_chain.ainvoke(message.content,
    #                                 callbacks=[
    #                                     cl.AsyncLangchainCallbackHandler()])

    response = await llm_chain.ainvoke({
        "chat_history": llm_chain.memory.load_memory_variables({})["chat_history"],
        "question": message.content
    }, callbacks=[
        cl.AsyncLangchainCallbackHandler()
    ])
    await cl.Message(response["text"]).send()


async def send_good_luck_message():
    await cl.Message(content="Good luck! 🍀", align="bottom").send()
  
async def handle_topic_selection(action: cl.Action):
    llm_chain = cl.user_session.get("llm_chain")
    #cl.user_session.set("selected_topic", action.value)
    #await cl.Message(content=f"Selected {action.value}").send()
    response = await llm_chain.ainvoke({
        "chat_history": llm_chain.memory.load_memory_variables({})["chat_history"],
        "question": f"Quiz me on the topic {action.value}."
    }, callbacks=[
        cl.AsyncLangchainCallbackHandler()
    ])
    await cl.Message(response["text"]).send()

@cl.action_callback("Linear Algebra")
async def on_action(action: cl.Action):
    await handle_topic_selection(action)

@cl.action_callback("Probability")
async def on_action(action: cl.Action):
    await handle_topic_selection(action)

@cl.action_callback("Accounts")
async def on_action(action: cl.Action):
    await handle_topic_selection(action)

@cl.action_callback("Calculus")
async def on_action(action: cl.Action):
    await handle_topic_selection(action)