import streamlit as st
from ultralytics import YOLO
from main_view import Main_View
from upload_view import Upload_View
from video_view import Video_View

# ✅ Initialize session state
if "page" not in st.session_state:
    st.session_state.page = "Main"  # Default page

if 'video_processed' not in st.session_state:
    st.session_state.video_processed = False


class App:
    def __init__(self):
        self.model = self.load_model()
        self.main_view = Main_View(self)
        self.upload_view = Upload_View(self, self.model)
        self.camera_view = Video_View(self, self.model)


    @st.cache_resource
    def load_model(_self):
        """Load the YOLOv8 model (cached)."""
        with st.spinner("Loading model... ⏳"):
            model = YOLO("best.pt")
        return model

    def change_page(self, page_name):
        """Change active page."""
        st.session_state.page = page_name
        st.rerun()  # Force rerun Streamlit to update UI

    def show(self):
        """Display only the active view."""
        if st.session_state.page == "Main":
            self.main_view.show()
        elif st.session_state.page == "Upload":
            self.upload_view.show()
        elif st.session_state.page == "Camera":
            self.camera_view.show()


# Start the app
app = App()
app.show()