from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from prompts import coding_hints_prompt_template
from langchain.memory.buffer import ConversationBufferMemory

from dotenv import load_dotenv

import chainlit as cl

# load_dotenv()


@cl.on_chat_start
def query_llm():
    openai_api_key = os.getenv("OPENAI_API_KEY")
    if not openai_api_key:
        raise ValueError("OPENAI_API_KEY environment variable not set")
        
    llm = ChatOpenAI(model="gpt-4o",
                     temperature=0.3)
    conversation_memory = ConversationBufferMemory(memory_key="chat_history",
                                                   return_messages=True,
                                                   )
    llm_chain = LLMChain(llm=llm, prompt=coding_hints_prompt_template, memory=conversation_memory)
    cl.user_session.set("llm_chain", llm_chain)


@cl.on_message
async def query_llm(message: cl.Message):
    llm_chain = cl.user_session.get("llm_chain")
    response = await llm_chain.acall(message.content,
                                     callbacks=[
                                         cl.AsyncLangchainCallbackHandler()])

    await cl.Message(response["text"]).send()

example_input = """Implement Selection Sort
Given a list of numbers, sort it using the Selection Sort algorithm.

Example
{
"arr": [5, 8, 3, 9, 4, 1, 7]
}
Output:

[1, 3, 4, 5, 7, 8, 9]


Constraints:
1 <= length of the given list <= 4 * 10^3
-10^9 <= number in the list <= 10^9"""