import pandas as pd
from os import environ
import datetime
import streamlit as st
from langchain.schema import HumanMessage, FunctionMessage

from helper import build_agents
from login import back_to_main

environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']

def on_chat_submit():
    ret = st.session_state.agents[st.session_state.sel][st.session_state.ret_type]({"input": st.session_state.chat_input})
    print(ret)
    
def clear_history():
    st.session_state.agents[st.session_state.sel][st.session_state.ret_type].memory.clear()


def back_to_main():
    if "user_info" in st.session_state:
        del st.session_state.user_info
    if "user_name" in st.session_state:
        del st.session_state.user_name
    if "jump_query_ask" in st.session_state:
        del st.session_state.jump_query_ask

def chat_page():
    st.session_state["agents"] = build_agents(f"{st.session_state.user_name}?default")
    with st.sidebar:
        st.radio("Retriever Type", ["Self-querying retriever", "Vector SQL"], key="ret_type")
        st.selectbox("Knowledge Base", ["ArXiv Papers", "Wikipedia", "ArXiv + Wikipedia"], key="sel")
        st.button("Clear Chat History", on_click=clear_history)
        st.button("Logout", on_click=back_to_main)
    for msg in st.session_state.agents[st.session_state.sel][st.session_state.ret_type].memory.chat_memory.messages:
        speaker = "user" if isinstance(msg, HumanMessage) else "assistant"
        if isinstance(msg, FunctionMessage):
            with st.chat_message("Knowledge Base", avatar="📖"):
                print(type(msg.content))
                st.write(f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*")
                st.write("Retrieved from knowledge base:")
                try:
                    st.dataframe(pd.DataFrame.from_records(map(dict, eval(msg.content))))
                except:
                    st.write(msg.content)
        else:
            if len(msg.content) > 0:
                with st.chat_message(speaker):
                    print(type(msg), msg.dict())
                    st.write(f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*")
                    st.write(f"{msg.content}")
    st.chat_input("Input Message", on_submit=on_chat_submit, key="chat_input")