import streamlit as st
import re
import streamlit_authenticator as stauth
import yaml
from yaml.loader import SafeLoader
from streamlit_player import st_player
from utils import create_transcript_from_youtube_api, create_open_ai_query
from prompts import DETECT_INTENT_OF_CONVERSATION, TOPIC_BASED_QUESTION, FOLLOW_UP_QUESTION, GENERAL_QUESTION, GENERAL_GREETING, \
    VAGUE_QUERY_PROMPT

st.set_page_config(page_title="Youtube AI")
with open('config.yaml') as file:
    config = yaml.load(file, Loader=SafeLoader)

authenticator = stauth.Authenticate(
    config['credentials'],
    config['cookie']['name'],
    config['cookie']['key'],
    config['cookie']['expiry_days'],
    config['preauthorized']
)
name, authentication_status, username = authenticator.login()

if st.session_state["authentication_status"]:
    authenticator.logout('Logout', 'main')
    st.write(f'Welcome to Mentor Mode')
elif st.session_state["authentication_status"] is False:
    st.error('Wrong password or username')
elif st.session_state["authentication_status"] is None:
    st.warning('Please enter your username and password')
    st.session_state["chat_history"] = []

if st.session_state["authentication_status"]:
    if "chat_history" not in st.session_state:
        st.session_state["chat_history"] = []

    if "messages" not in st.session_state:
        st.session_state.messages = []

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

    with st.sidebar:
        st.title("Your Video")
        youtube_video_link = st.text_area("Please enter your video link")
        st.button("Play Video", type="primary")
        if youtube_video_link:
            st_player(youtube_video_link)
        else:
            st.write("Please enter a valid link")

    if prompt := st.chat_input("Hey AI!"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.session_state.chat_history.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

    if youtube_video_link and prompt:
        with st.spinner("Processing..."):
            video_id = re.search(r'(?<=v=)[\w-]+', youtube_video_link).group(0)
            yt_transcript = create_transcript_from_youtube_api(video_id)
            if yt_transcript["success"]:
                ADDITIONAL_PROMPT = f"""QUERY : ```{prompt}```, TRANSCRIPT:```{yt_transcript}```, 
                CHAT_HISTORY:```{st.session_state["chat_history"]}````"""
                FINAL_PROMPT = ADDITIONAL_PROMPT + DETECT_INTENT_OF_CONVERSATION
                intent = create_open_ai_query(FINAL_PROMPT)
                print(intent["data"])
                if intent["success"]:
                    if intent["data"] == "VAGUE_QUERY":
                        FINAL_PROMPT = ADDITIONAL_PROMPT + VAGUE_QUERY_PROMPT
                        response = create_open_ai_query(FINAL_PROMPT)
                    elif intent["data"] == "GENERAL_QUESTION":
                        FINAL_PROMPT = ADDITIONAL_PROMPT + GENERAL_QUESTION
                        response = create_open_ai_query(FINAL_PROMPT)
                    elif intent["data"] == "TOPIC_BASED_QUESTION":
                        FINAL_PROMPT = ADDITIONAL_PROMPT + TOPIC_BASED_QUESTION
                        response = create_open_ai_query(FINAL_PROMPT)
                    elif intent["data"] == "FOLLOW_UP_QUESTION":
                        FINAL_PROMPT = ADDITIONAL_PROMPT + FOLLOW_UP_QUESTION
                        response = create_open_ai_query(FINAL_PROMPT)
                    elif intent["data"] == "GENERAL_GREETING":
                        FINAL_PROMPT = ADDITIONAL_PROMPT + GENERAL_GREETING
                        response = create_open_ai_query(FINAL_PROMPT)

        with st.chat_message("assistant"):
            if response["success"]:
                st.write(response["data"])
            else:
                st.write(response["error"])
        st.session_state.messages.append({"role": "assistant", "content": response["data"]})
        st.session_state.chat_history.append({"role": "assistant", "content": response["data"]})

    if st.button("Download Chat History"):
        # Combine role and content for each message
        chat_history = "\n".join(
                [f"{message['role']} : {message['content']}" for message in
                 st.session_state.messages if message["content"] is not None]
                )
        st.download_button(
                label="Download",
                data=chat_history,
                file_name="chat_history.txt",
                mime="text/plain"
        )