import chainlit as cl
import os
from dotenv import load_dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables.config import RunnableConfig
from langchain_openai import ChatOpenAI

load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")

role = ""
attitude = ""

@cl.action_callback("ceo")
async def on_action(action):
    role = "ceo"
    cl.user_session.set("role", role)
    await cl.Message(content=f"The CEO has entered the building").send()

@cl.action_callback("vp-sales")
async def on_action(action):
    role = "vp-sales"
    cl.user_session.set("role", role)
    await cl.Message(content="The VP of Sales has entered the building").send()

@cl.action_callback("happy")
async def on_action(action):
    attitude = "happy but hates being called sandy"
    cl.user_session.set("attitude", attitude)
    role = cl.user_session.get("role", "ceo")
    message = f"The {role} is happy"
    await cl.Message(content=message).send()

@cl.action_callback("angry")
async def on_action(action):
    attitude = "angry but could be convined to calm down"
    cl.user_session.set("attitude", attitude)
    role = cl.user_session.get("role", "ceo")
    message = f"The {role} is angry"
    await cl.Message(message).send()

@cl.action_callback("sarcastic")
async def on_action(action):
    personality = "sarcastic"
    cl.user_session.set("personality", personality)
    role = cl.user_session.get("role", "ceo")
    message = f"The {role} is sarcastic"
    await cl.Message(message).send()

@cl.action_callback("forgetful")
async def on_action(action):
    personality = "forgetful"
    cl.user_session.set("personality", personality)
    role = cl.user_session.get("role", "ceo")
    message = f"The {role} is forgetful"
    await cl.Message(message).send()

@cl.action_callback("attentitive")
async def on_action(action):
    personality = "attentitive"
    cl.user_session.set("personality", personality)
    role = cl.user_session.get("role", "ceo")
    message = f"The {role} is attentitive"
    await cl.Message(message).send()

user_template = """
        Conversation History:
        {conversation_history}

        Question:
        {question}

        Company:
        {company}

        Role:
        {role}

        Attitude:
        {attitude}

        Personality:
        {personality}
        Twist:
        {twist}


    """

system_template = """
    You playing a role in a conversation with a sales representative.
    You are a customer of his.
    You are asking him questions about the product.
    You role is defined by the {role} section.
    Your attitude is defined by the {attitude} section.
    There may be a twist that you should keep in mind. This is found in the {twist} section.
    The twist will not always be present.
""" 

#############################################
### On Chat Start (Session Start) Section ###
#############################################
@cl.on_chat_start
async def on_chat_start():
    # create a chain

    company = "Acme Corp"
    cl.user_session.set("company", company)
    cl.user_session.set("message_count", 0)
    twist = ""
    cl.user_session.set("twist", twist)



    chat_prompt = ChatPromptTemplate.from_messages([
        ("system", system_template),
        ("human", user_template)
    ])
    chat_model = ChatOpenAI(model="gpt-4o-mini")
    simple_chain = chat_prompt | chat_model
    cl.user_session.set("chain", simple_chain)

    role_actions = [
        cl.Action(name="ceo", value="ceo", description="CEO"),
        cl.Action(name="vp-sales", value="vp-sales", description="VP of Sales")
    ]
    await cl.Message(content="Select a role", actions=role_actions).send()

    
 
    attitude_actions = [
        cl.Action(name="angry", value="angry", description="Angry"),
        cl.Action(name="happy", value="happy", description="Happy")
    ]

    await cl.Message(content="Select an attitude", actions=attitude_actions).send()

    personality_actions = [
        cl.Action(name="sarcastic", value="sarcastic", description="Sarcastic"),
        cl.Action(name="forgetful", value="forgetful", description="Forgetful"),
        cl.Action(name="attentitive", value="attentitive", description="Attentitive")
    ]
    await cl.Message(content="Select a personality", actions=personality_actions).send()    

@cl.on_message
async def main(message: cl.Message):
    history = cl.user_session.get("history", [])
    message_count = cl.user_session.get("message_count", 0)
    message_count += 1
    cl.user_session.set("message_count", message_count)
    chain = cl.user_session.get("chain")
    company = cl.user_session.get("company", "Acme Corp")
    role = cl.user_session.get("role", "ceo")
    attitude = cl.user_session.get("attitude", "happy")
    personality = cl.user_session.get("personality", "sarcastic")
    twist = cl.user_session.get("twist", "")
    question = message.content
    if message_count == 2:
        twist = "The fire alarm is going off"
    msg = cl.Message(content="")

    history.append({"role": "user", "content": message})
    # handle streaming of LLM responses 
    response_content = ""

    async for chunk in chain.astream(
        {"question": question, "company": company, "role": role, "attitude": attitude, "twist": twist, "conversation_history": history, "personality": personality},
        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
    ):
        response_content += chunk.content
        await msg.stream_token(chunk.content)
    history.append({"role": "assistant", "content": response_content})
    cl.user_session.set("history", history)
    await msg.send()