from dotenv import load_dotenv
import streamlit as st
from receptionist_agent import get_agent_response
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage, messages_to_dict
import json
from datetime import datetime
import time
from download import format_conversation_for_download

#import langsmith

load_dotenv()

#langsmith.debug = True

st.set_page_config(layout="wide")

if "conversation" not in st.session_state:
    st.session_state.conversation = [AIMessage(content="Thanks for calling, this is Receptionist Agent! Can I help you schedule a service or would you like help with something else?")]
    start_time = datetime.now().strftime("%y%m%d-%H%M%S")
    st.session_state.conversation_file = f"conversations/conversation_{start_time}.json"

def save_conversation():
    with open(st.session_state.conversation_file, "w") as f:
        json.dump(messages_to_dict(st.session_state.conversation), f, indent=4)

def submit_message():
    message = st.session_state.message
    st.session_state.conversation.append(HumanMessage(content=message))
    
    start = time.time()
    new_messages = get_agent_response(st.session_state.conversation)
    end = time.time()
    st.session_state.response_time = f"{end - start:.2f}"
    
    st.session_state.conversation.extend(new_messages)
    # save_conversation()

st.subheader("Receptionist Agent")

for message in st.session_state.conversation:
    if isinstance(message, HumanMessage):
        with st.chat_message("user"):
            st.write(message.content)
    elif isinstance(message, AIMessage):
        with st.chat_message("assistant"):
            if message.content:
                st.write(message.content)
            if message.tool_calls:
                for tool_call in message.tool_calls:
                    st.write(f"Calling tool: `{tool_call['name']}` with args `{tool_call['args']}`")
    elif isinstance(message, ToolMessage):
        with st.chat_message("tool"):
            st.write(f"Tool `{message.tool_call_id}` output: `{message.content}`")

if len(st.session_state.conversation) > 1:
    _, col2 = st.columns([0.75, 0.25])
    with col2:
        st.download_button(
            label="Download Conversation",
            data=format_conversation_for_download(st.session_state.conversation),
            file_name=st.session_state.conversation_file.replace(".json", "_formatted.json"),
            mime="application/json",
        )

if "response_time" in st.session_state:
    st.caption(f"Response time: {st.session_state.response_time}s")

st.chat_input("Type message here", on_submit=submit_message, key="message")