from flask import session
from typing import Any


class StateManager:
    def __init__(self):
        self.should_rerun = False
        self.form_data = {}
    
    def init_session(self):
        """Initialize session state"""
        if 'feather_state' not in session:
            session['feather_state'] = {}
    
    def set_form_data(self, data: dict):
        """Store form submission data"""
        self.form_data = data
    
    def get_form_data(self, key: str, default=None):
        """Get form data by key"""
        return self.form_data.get(key, default)
    
    def trigger_rerun(self):
        """Trigger a page rerun"""
        self.should_rerun = True


_state_manager = StateManager()


def rerun():
    """Trigger a page rerun"""
    _state_manager.trigger_rerun()


def set_state(key: str, value: Any):
    """Set a session state variable"""
    if 'feather_state' not in session:
        session['feather_state'] = {}
    session['feather_state'][key] = value
    session.modified = True


def get_state(key: str, default=None) -> Any:
    """Get a session state variable"""
    if 'feather_state' not in session:
        return default
    return session['feather_state'].get(key, default)
