import streamlit as st
from time import sleep

from src.st_helpers import st_setup
from src.common import *
from src.architectures import Architecture
from src.testing import TestGroup


if st_setup('LLM Arch'):
    st.write("# System Status")

    if 'admin_logged_in' not in st.session_state:
        entered_pw = st.text_input(label="Enter the admin password to manage the system", type="password")
        if st.button('Login'):
            if entered_pw == st.secrets['admin_pw']:
                st.session_state['admin_logged_in'] = True
                st.rerun()
            else:
                st.error("Incorrect password")
    else:
        st.write("## Wipe Trace Logs")
        if st.button("Wipe logs"):
            Architecture.wipe_trace()
        st.write('Note - wipe will only wipe the temporary file, DB persisted records will be saved')

        st.divider()

        if st.button("Reload traces"):
            TestGroup.load_all(True)

        st.divider()

        st.write("## HF Inference Endpoint Statuses")
        st.write("The following endpoints need to be running to run all the demonstrations.")

        endpoints = st.secrets['endpoints'].split(',')
        refresh = False

        for i, e in enumerate(endpoints):
            status = hf_endpoint_status('alfraser', e)
            message = f'{e} ({status})'

            if i != 0:
                st.divider()

            status_col, button_col = st.columns([2, 1])

            if status == HF_RUNNING:
                with status_col:
                    st.success(message)
                with button_col:
                    if st.button("Pause", key=f'action_{i}'):
                        pause_hf_endpoint('alfraser', e)
                        st.rerun()
            elif status == HF_SCALEDTOZERO:
                with status_col:
                    st.error(message)
            elif status == HF_PAUSED:
                with status_col:
                    st.warning(message)
                with button_col:
                    if st.button("Resume", key=f'action_{i}'):
                        resume_hf_endpoint('alfraser', e)
                        st.rerun()
            elif status == HF_FAILED:
                with status_col:
                    st.error(message)
                with button_col:
                    if st.button("Pause (to restart)", key=f'action_{i}'):
                        pause_hf_endpoint('alfraser', e)
                        st.rerun()
            else:
                refresh = True
                with status_col:
                    st.info(message)
                with button_col:
                    if st.button("Pause", key=f'action_{i}'):
                        pause_hf_endpoint('alfraser', e)
                        st.rerun()

        if refresh:
            with st.spinner('Updating every 10 seconds...'):
                sleep(10)
                st.rerun()