utils/session.py

import uuid

import streamlit as st

from config import AppConfig
from utils.data_loading import DataManager


def init_session_state(config: AppConfig) -> None:
    """Initialize session state with better separation of concerns"""
    # First initialize config-based states
    if "config_initialized" not in st.session_state:
        st.session_state.update(
            {
                "DEBUG": config.DEBUG,
                "ENABLE_TIMING": config.ENABLE_TIMING,
                "ENABLE_ALTAIR": config.ENABLE_ALTAIR,
                "reporting_month": config.DEFAULT_REPORTING_MONTH,
            }
        )
        st.session_state.config_initialized = True

    # Initialize analytics state
    if "analytics_initialized" not in st.session_state:
        st.session_state.update(
            {
                "logged_visit": False,
                "visitor_id": str(uuid.uuid4()),
            }
        )
        st.session_state.analytics_initialized = True

    # Initialize admin state
    if "admin_initialized" not in st.session_state:
        st.session_state.update(
            {
                "admin_authenticated": False,
            }
        )
        st.session_state.admin_initialized = True

    # Initialize exclusion settings
    if "exclusions_initialized" not in st.session_state:
        # Start with default exclusions from config
        st.session_state.update(
            {
                "persistent_excluded_sectors": config.DEFAULT_EXCLUDED_SECTORS,
                "persistent_excluded_stations": config.DEFAULT_EXCLUDED_STATIONS,
                "exclusions_initialized": True,
            }
        )

    # Initialize data manager if not already initialized
    if "data_manager" not in st.session_state:
        st.session_state.data_manager = DataManager(config)

    # Ensure excluded values are valid using complete lists from data manager
    if "data_manager" in st.session_state:
        # Get valid options from data manager
        valid_sectors = set(st.session_state.data_manager.all_sectors)
        valid_stations = set(st.session_state.data_manager.all_stations)

        # Filter out invalid values
        if "persistent_excluded_sectors" in st.session_state:
            st.session_state.persistent_excluded_sectors = [
                sector
                for sector in st.session_state.persistent_excluded_sectors
                if sector in valid_sectors
            ]

        if "persistent_excluded_stations" in st.session_state:
            st.session_state.persistent_excluded_stations = [
                station
                for station in st.session_state.persistent_excluded_stations
                if station in valid_stations
            ]

    # Initialize data and date-related variables
    if "data" not in st.session_state:
        metadata = st.session_state.data_manager.metadata
        st.session_state.update(
            {
                "data": st.session_state.data_manager.load_data() if metadata else None,
                "start_date": config.DEFAULT_START_DATE,
                "end_date": config.DEFAULT_END_DATE,
                "reporting_month": config.DEFAULT_REPORTING_MONTH,
            }
        )


def ensure_session_initialized(config=None):
    """Ensures session state is properly initialized with all critical variables."""
    if "session_initialized" in st.session_state:
        return

    # If no config provided, try to load it
    if config is None:
        try:
            from config import AppConfig

            config = AppConfig.from_env()
        except Exception as e:
            st.error(f"Failed to load configuration: {str(e)}")
            st.stop()

    try:
        init_session_state(config)
        st.session_state.session_initialized = True
        if config.DEBUG:
            st.info(f"Session initialized with config: {config}")
    except Exception as e:
        st.error(f"Failed to initialize session state: {str(e)}")
        if config.DEBUG:
            st.exception(e)
        st.stop()

    # Verify critical session state variables are set
    critical_vars = [
        "data_manager",
        "config_initialized",
        "analytics_initialized",
        "admin_initialized",
    ]
    missing_vars = [var for var in critical_vars if var not in st.session_state]
    if missing_vars:
        st.error(f"Critical session state variables missing: {missing_vars}")
        st.stop()