from langchain.schema import SystemMessage, HumanMessage, AIMessage
from memory import memory
from agent import agent


def respond(message, history, system_message, max_tokens, temperature, top_p):
    memory.chat_memory.add_message(SystemMessage(content=system_message))
    for user_input, bot_response in history:
        if user_input:
            memory.chat_memory.add_message(HumanMessage(content=user_input))
        if bot_response:
            memory.chat_memory.add_message(AIMessage(content=bot_response))
    memory.chat_memory.add_message(HumanMessage(content=message))

    if "flight" in message.lower():
        message = f"Use the 'Flight Booking' tool. {message}"
    elif "hotel" in message.lower() or "stay" in message.lower():
        message = f"Use the 'Hotel Booking' tool. {message}"

    response_dict = agent.invoke({"input": message})  # Updated from .run to .invoke
    response = response_dict["output"] if "output" in response_dict else str(response_dict)

    if "❌" in response:
        response += "\n\nWould you like me to retry with different dates (e.g., 'hotels in Paris from March 1 2025 to March 4 2025') or assist with something else (e.g., flights)?"
    return response