import time
from os.path import basename
import yaml

import streamlit as st
from yaml.loader import SafeLoader
from streamlit_authenticator import Authenticate

from htmlTemplates import css, bot_template, user_template
from sources import sources_ref
import constants
import backend


import torch
print(f"Is CUDA available: {torch.cuda.is_available()}")
# True
if torch.cuda.is_available():
    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
# Tesla T4


def on_query():
    constants.question = st.session_state.InputText
    st.session_state.InputText = ""


def handle_userinput(user_question):
    res = st.session_state.model({'query': user_question})
    answer, docs = res["result"], res["source_documents"]
    # source = sources_ref[basename(docs[0].metadata["source"])]

    st.session_state.chat_history.insert(0, user_question)
    # st.session_state.chat_history.insert(1, answer + "<br><br><em>Source: " + source)
    st.session_state.chat_history.insert(1, answer)

    for i, message in enumerate(st.session_state.chat_history):
        if i % 2 == 0:
            st.write(user_template.replace("{{MSG}}", message), unsafe_allow_html=True)
        else:
            st.write(bot_template.replace("{{MSG}}", message), unsafe_allow_html=True)


def login_page():
    st.set_page_config(page_title="logotherapyGPT",
                       page_icon=":books:")

    with open('auth_config.yaml') as file:
        config = yaml.load(file, Loader=SafeLoader)

    authenticator = Authenticate(
        config['credentials'],
        config['cookie']['name'],
        config['cookie']['key'],
        config['cookie']['expiry_days'],
        config['preauthorized']
    )

    name, authentication_status, username = authenticator.login('Login', 'main')

    if authentication_status:

        if 'db_clicked' not in st.session_state:
            st.session_state.db_clicked = False

        placeholder_box = st.empty()
        placeholder_button = st.empty()
        placeholder_box.selectbox(
            'Choose a Database',
            ("All", "Frankl's Works", "Journal of Search for Meaning"),
            key="database"
        )
        placeholder_button.button(
            "Enter",
            on_click=db_button_clicked
        )

        if st.session_state.db_clicked:
            placeholder_box.empty()
            placeholder_button.empty()
            main_page()

    elif not authentication_status:
        st.error('Username/password is incorrect')

    elif authentication_status is None:
        st.warning('Please enter your username and password')


def db_button_clicked():
    st.session_state.db_clicked = True


def main_page():
    st.write(css, unsafe_allow_html=True)

    st.markdown("# logotherapyGPT")
    st.markdown("Chosen Database: " + st.session_state.database)
    st.markdown("An AI chatbot to quickly access logotherapy information. Trained on the works of Viktor "
                "Frankl and The International Forum for Logotherapy.")

    if "model" not in st.session_state:
        with st.spinner("Loading Database..."):
            st.session_state.model = backend.load_qa(st.session_state.database)

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

    st.text_input(
        "Input",
        label_visibility="collapsed",
        placeholder="Ask a Question...",
        on_change=on_query,
        key="InputText"
    )

    print(constants.question)

    if st.session_state.model is not None and constants.question is not None:
        start = time.time()

        with st.spinner("Processing..."):
            handle_userinput(constants.question)

        end = time.time()
        print(f"\n> Answer (took {round(end - start, 2)} s.):")


if __name__ == '__main__':
    login_page()