import random
import streamlit as st
from taai_demo.ui.access_page import access_page
from taai_demo.ui.chat_page import chat_page
from taai_demo.ui.search_page import search_page
from taai_demo.ui.upload_page import upload_page
from taai_demo.ui.state_management import transition_state
from taai_demo.settings import Settings


def initialize_state(name, default):
    if name not in st.session_state:
        st.session_state[name] = default


def render_state():
    states = {
        "access": access_page,
        "chat": chat_page,
        "search": search_page,
        "upload": upload_page,
    }
    states[st.session_state.state]()


def main():
    initialize_state("authorized", False)
    initialize_state("state", "access")
    initialize_state("chat_history", [])
    initialize_state("semantic_search_wrapper", None)
    initialize_state("settings", Settings())
    initialize_state("message_seed", random.randint(0, 1e3))

    if st.session_state.state in ("chat", "search"):
        new_state = st.sidebar.selectbox("Choose a demo", ["Search", "Chat"])
        transition_state(new_state.lower())
        if st.button("Reupload File"):
            transition_state("upload")

    render_state()


if __name__ == "__main__":
    main()