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

from pages import landing_page, model_inference_page

if "page" not in sst:
    sst['page'] = 'landing_page'

def reset_sst():
    for key in list(sst.keys()):
        if key != "page" and key != 'device':
            sst.pop(key, None)


# 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"] == "landing_page":
        reset_sst() # reset all session state variables before navigating to the landing page
        await landing_page()  # Call the landing page function

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

asyncio.run(main())