import streamlit as st
import os
import warnings
from langchain.agents import initialize_agent, AgentType
from chat_agent import llm, order_query_tool, answer_tool
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage

warnings.filterwarnings("ignore")

st.title("🍔 FoodDelivery Chat Bot Agent....")

# 1. Initialize variables ONLY if they don't exist
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

if "agent" not in st.session_state:
    tools = [order_query_tool, answer_tool]
    st.session_state.agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
        verbose=True, # Set to True for debugging
        handle_parsing_errors=True
    )

# 2. Display existing history
for message in st.session_state.chat_history:
    if isinstance(message, HumanMessage):
        st.chat_message("user").write(message.content)
    elif isinstance(message, AIMessage):
        st.chat_message("assistant").write(message.content)

cust_id = 'C1014'

# 3. Chat Logic
if prompt := st.chat_input("Ask your query here...."):
    # Pre-process the prompt
    full_prompt = f"{prompt} where cust_id is {cust_id}"
    
    # Display the user message immediately
    st.chat_message("user").write(prompt)
    
    with st.spinner("Searching..."):
        try:
            # LangChain Agent call
            response = st.session_state.agent.invoke({
                "input": full_prompt,
                "chat_history": st.session_state.chat_history
            })
            
            answer = response["output"]
            
            # 4. Update Chat History properly
            st.session_state.chat_history.append(HumanMessage(content=full_prompt))
            st.session_state.chat_history.append(AIMessage(content=answer))
            
            # Display Assistant response
            st.chat_message("assistant").write(answer)
            
        except Exception as e:
            st.error(f"An error occurred: {e}")