import streamlit as st
import time

from src.st_helpers import st_setup
from src.common import *


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

    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()
        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...'):
            time.sleep(5)
            st.rerun()