import streamlit as st
from YoutubeChat.YoutubeChat import YoutubeChatBot
from openai import AuthenticationError

def get_user_input():
    st.title("Setup")

    # Input field for YouTube link
    youtube_link = st.text_input("Enter YouTube Video Link")

    # Input field for OpenAI API key
    openai_api_key = st.text_input("Enter OpenAI API Key", type = "password", help="Please delete the API key after use for your own security.")

    # Button to start chat
    if st.button("Start Chat"):
        if youtube_link.strip() == "":
            st.error("Please enter a valid YouTube link.")
        elif openai_api_key.strip() == "":
            st.error("Please enter your OpenAI API key.")
        else:
            try:
                yt = YoutubeChatBot(youtube_link, OPENAI_API_KEY=openai_api_key)
                start_message = yt.send('start protocol')
                st.session_state.start_message = start_message
                st.session_state.chat_object = yt
                st.session_state.messages = [{'role':'assistant', 'content': start_message}]
                st.rerun()  # Rerun the app to switch to the main chat page
            except AuthenticationError as e:
                st.error("Incorrect API key provided. Please check your API key.")

                # Optionally provide guidance on how to resolve the issue
                st.write("You can find your API key at https://platform.openai.com/account/api-keys.")
            except ValueError as e:
                st.error("Incorrect video link provided! ")


def main_chat_page():
    st.title("Main Chat Page")

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

    # Display chat input box
    prompt = st.chat_input("Ask something")
    
    # If user inputs something
    if prompt:
        st.session_state.messages.append({"role": "user", "content": prompt})

        with st.chat_message("user"):
            st.write(prompt)
        
        result = st.session_state.chat_object.send(prompt)
        st.session_state.messages.append({'role':'assistant', "content": result})

        with st.chat_message("assistant"):
            st.write(result)

if __name__ == "__main__":
    if "start_message" not in st.session_state:
        get_user_input()
    else:
        main_chat_page()