import os
import time
import json
import streamlit as st
from google_api_deploy import read_google_doc, write_google_doc_add, authenticate_google_drive, write_google_doc, monitor_google_doc




service = authenticate_google_drive()

# Streamlit page configuration
st.set_page_config(
    page_title="AMGPT",
    page_icon="🦙",
    layout="centered",
    initial_sidebar_state="auto",
    menu_items=None
)
st.title("AMGPT, powered by LlamaIndex 💬🦙")
st.info("Please hit the refresh button just once before prompting")

# Initialize chat messages history
if "messages" not in st.session_state.keys():
    st.session_state.messages = [
        {"role": "assistant", "content": "Ask me a question!!!"}
    ]


if st.sidebar.button("Refresh"):
    st.session_state.messages = [{"role": "assistant", "content": "Ask me a question!!!"}]
    write_google_doc(service, "1PMgaUKuPXYrbvYiqx7lwIG47P_wzxMN7VmSxg1RMIKs", '')#prompt
    write_google_doc(service, "1Na3EF7vWj9x_yzVxs3w450zgWtIdZZHU0xOYV-9WPk4", '')#promptandquery
    write_google_doc(service, "1Wpscos9GJ04_T6eLWHbYq2f6HuZD3OJvyoqW_pE2Yoc", '')#query


# Prompt for user input and save to chat history
if prompt := st.chat_input("Your question"):
    with st.spinner("Thinking..."):
        write_google_doc_add(service, "1Na3EF7vWj9x_yzVxs3w450zgWtIdZZHU0xOYV-9WPk4", prompt)#promptandquery
        write_google_doc_add(service, "1PMgaUKuPXYrbvYiqx7lwIG47P_wzxMN7VmSxg1RMIKs", prompt)#prompt
        st.session_state.messages.append({"role": "user", "content": prompt})
        response = monitor_google_doc(service, "1Wpscos9GJ04_T6eLWHbYq2f6HuZD3OJvyoqW_pE2Yoc")#reading query
        

for message in st.session_state.messages[1:]:
    with st.chat_message(message["role"]):
        st.write(message["content"])




if st.session_state.messages[-1]["role"] != "assistant":
    with st.chat_message("assistant"):
        st.write(response)
        message = {"role": "assistant", "content": response}
        # Add response to message history
        st.session_state.messages.append(message)



# # Generate a new response if the last message is not from the assistant 
# if st.session_state.messages[-1]["role"] != "assistant":     
#         with st.chat_message("assistant"):
#                 history = [msg["content"] for msg in st.session_state.messages]
#                 total_history = " ".join(history) + " " + st.session_state["additional_string"]
#                 #response = st.session_state.chat_engine.chat(st.session_state.messages[-1]["content"])
#                 response = st.session_state.chat_engine.chat(total_history)
#                 if response.response == "Empty Response":
#                     response = st.session_state.llm.complete(total_history)
#                     st.write(response.text)
#                     message = {"role": "assistant", "content": response.text}
#                 else:
#                     st.write(response.response)
#                     message = {"role": "assistant", "content": response.response}
#                 st.session_state.messages.append(message)