from utils import navigate_to
from pages import manual_input_page, image_input_page, model_inference_page

import streamlit as st
from streamlit import session_state as sst
import asyncio

#TODO: Fix model inference and post processing function befor emoving ot production.
# Initialize session state variable to start with home page
if "page" not in sst:
    sst["page"] = "Home"

# function to remove all sesion variables from sst, except page.
def reset_sst():
    for key in list(sst.keys()):
        if key != "page":
            sst.pop(key, None)

# Landing page function
async def landing_page():
    
    st.title("We will explain your menu like never before!")
    st.write("\n")
    st.write("\n")
    st.write("\n")
    
    c1, c2= st.columns(2)

    with c1:
        # Navigate to manual input page if user clicks on the button
        st.button("Enter Items Manually", on_click=navigate_to, args=("ManualInput",))

    with c2:
        # Navigate to image input page if user clicks on the button
        st.button("Upload Items from Image", on_click=navigate_to, args=("ImageInput",))

    

# Main function to handle navigation
async def main():
    """
    Main function that handles the navigation logic based on the current page.

    Returns:
        None
    """

    # Navigation logic
    if sst["page"] == "Home":
        reset_sst() # reset all session state variables before navigating to the landing page
        await landing_page()  # Call the landing page function
            
    elif sst["page"] == "ManualInput":
        reset_sst() # reset all session state variables before navigating to the landing page
        await manual_input_page()  # Call the manual input page function
    
    elif sst["page"] == "ImageInput":
        reset_sst() # reset all session state variables before navigating to the landing page
        await image_input_page()  # Call the image input page function

    elif sst["page"] == "Inference":
        await model_inference_page()  # Call the model inference page function

asyncio.run(main())