import streamlit as st

from src.common import *


def login(pw: str) -> bool:
    """
    Convenience wrapper for login functionality.  Has side effect of storing the logged in
    user name in the session state
    :param pw: Password to attempt login
    :return: True/False if logged in
    """
    pw_users = st.secrets['app_password'].split(';')
    for pu in pw_users:
        split_point=pu.find('(')
        pu_password = pu[0:split_point]
        if pu_password == pw:
            pu_user = pu[split_point + 1: -1]
            st.session_state['username'] = pu_user
            return True
    return False


def show_system_status_bar() -> None:
    endpoints = st.secrets['endpoints'].split(',')
    num_endpoints = len(endpoints)
    running = []
    stopped = []
    for e in endpoints:
        status = hf_endpoint_status('alfraser', e)
        if status == HF_RUNNING:
            running.append(e)
        else:
            stopped.append(f'{e} ({status})')

    if len(running) == num_endpoints:
        message = "**:large_green_circle: System Status - All LLM Endpoints Running**"
    elif len(stopped) == num_endpoints:
        message = "**:red_circle: System Status - No LLM Endpoints Running**"
    else:
        message = "**:large_orange_circle: System Status - Some LLM Endpoints Running**"

    with st.expander(message):
        if len(running) == num_endpoints:
            st.success("All LLM endpoints running")
            st.write("**Note** - endpoints shutdown after 15 minutes of inactivity")
        else:
            st.info(f"{len(running)} of {num_endpoints} LLM endpoints running")
            all = []
            for e in running:
                all.append(f'* :large_green_circle: {e}')
            for e in stopped:
                all.append(f'* :red_circle: {e}')
            st.write('\n'.join(all))
            if st.button("Start all endpoints"):
                for e in stopped:
                    resume_hf_endpoint('alfraser', e)
                st.write(
                    '**Note** - endpoint start may take multiple minutes, so refresh and check status periodically')


def st_setup(page_title: str, layout: str = 'wide', skip_login: bool = False) -> None:
    """
    Sets up standard outline (wide layout), checks for logged in status and then
    displays the login box if not logged in.  Should be used as a conditional to display
    the other content on the page.
    :param page_title: The streamlit page title for this page
    :return: bool - logged in status
    """
    st.set_page_config(page_title=page_title, layout=layout)

    username='username'
    if username in st.session_state:
        st.markdown(f'<div style="float: right; padding-top: 10px;">Logged in as <b>{st.session_state[username]}</b></span>', unsafe_allow_html=True)

    with st.sidebar:
        st.image('img/uob-logo.png', width=200)

    if skip_login:
        show_system_status_bar()
        return True

    # Option to add a running_local key to secrets.toml for local login to speed up development
    # This relies on the protection of streamlit secrets in production, which password login already relies on
    running_local = 'running_local'
    if running_local in st.secrets:
        show_system_status_bar()
        return True

    logged_in = 'logged_in'
    if logged_in not in st.session_state or st.session_state[logged_in]==False:
        _, c2, _ = st.columns(3)
        with c2:
            st.write('### Log in')
            pw = st.text_input(type='password', label='Password')
            if st.button('Log in'):
                if login(pw):
                    st.session_state[logged_in]=True
                    st.rerun()
                else:
                    st.info("Invalid password")
        return False
    else:
        show_system_status_bar()
        return True