utils/state.py

import uuid
import time
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Union
import json

# Import utilities
from utils.logging import get_logger
from utils.error_handling import handle_exceptions, ValidationError, safe_get
from utils.storage import save_data, load_data

logger = get_logger(__name__)

# Global state storage
_state_storage = {}
_activity_log = []

@handle_exceptions
def generate_id(prefix: str = "") -> str:
    """
    Generate a unique ID
    
    Args:
        prefix: Optional prefix for the ID
    
    Returns:
        str: Unique identifier
    """
    unique_id = str(uuid.uuid4())[:8]
    if prefix:
        return f"{prefix}_{unique_id}"
    return unique_id

@handle_exceptions
def get_timestamp(format_type: str = "iso") -> str:
    """
    Get current timestamp
    
    Args:
        format_type: Format type ('iso', 'unix', 'readable')
    
    Returns:
        str: Formatted timestamp
    """
    now = datetime.now(timezone.utc)
    
    if format_type == "iso":
        return now.isoformat()
    elif format_type == "unix":
        return str(int(now.timestamp()))
    elif format_type == "readable":
        return now.strftime("%Y-%m-%d %H:%M:%S UTC")
    else:
        return now.isoformat()

@handle_exceptions
def record_activity(activity_type: str, details: Dict[str, Any] = None) -> bool:
    """
    Record an activity in the activity log
    
    Args:
        activity_type: Type of activity
        details: Additional details about the activity
    
    Returns:
        bool: True if recorded successfully
    """
    try:
        activity = {
            "id": generate_id("activity"),
            "type": activity_type,
            "timestamp": get_timestamp(),
            "details": details or {}
        }
        
        _activity_log.append(activity)
        
        # Keep only last 1000 activities
        if len(_activity_log) > 1000:
            _activity_log.pop(0)
        
        logger.info(f"Activity recorded: {activity_type}")
        return True
        
    except Exception as e:
        logger.error(f"Failed to record activity: {e}")
        return False

@handle_exceptions
def get_state(key: str, default: Any = None) -> Any:
    """
    Get value from global state
    
    Args:
        key: State key
        default: Default value if key not found
    
    Returns:
        Any: State value or default
    """
    return safe_get(_state_storage, key, default)

@handle_exceptions
def set_state(key: str, value: Any) -> bool:
    """
    Set value in global state
    
    Args:
        key: State key
        value: Value to set
    
    Returns:
        bool: True if set successfully
    """
    try:
        _state_storage[key] = value
        logger.debug(f"State set: {key}")
        return True
    except Exception as e:
        logger.error(f"Failed to set state {key}: {e}")
        return False

@handle_exceptions
def update_state(key: str, updates: Dict[str, Any]) -> bool:
    """
    Update nested state values
    
    Args:
        key: State key
        updates: Dictionary of updates to apply
    
    Returns:
        bool: True if updated successfully
    """
    try:
        current = get_state(key, {})
        if isinstance(current, dict):
            current.update(updates)
            return set_state(key, current)
        else:
            logger.warning(f"Cannot update non-dict state: {key}")
            return False
    except Exception as e:
        logger.error(f"Failed to update state {key}: {e}")
        return False

@handle_exceptions
def clear_state(key: Optional[str] = None) -> bool:
    """
    Clear state (specific key or all)
    
    Args:
        key: Specific key to clear, or None to clear all
    
    Returns:
        bool: True if cleared successfully
    """
    try:
        if key:
            if key in _state_storage:
                del _state_storage[key]
                logger.info(f"State cleared: {key}")
        else:
            _state_storage.clear()
            logger.info("All state cleared")
        return True
    except Exception as e:
        logger.error(f"Failed to clear state: {e}")
        return False

@handle_exceptions
def get_activity_log(limit: int = 100) -> List[Dict[str, Any]]:
    """
    Get recent activity log entries
    
    Args:
        limit: Maximum number of entries to return
    
    Returns:
        List[Dict]: Activity log entries
    """
    try:
        return _activity_log[-limit:] if _activity_log else []
    except Exception as e:
        logger.error(f"Failed to get activity log: {e}")
        return []

@handle_exceptions
def save_state_to_file(filename: str = "app_state.json") -> bool:
    """
    Save current state to file
    
    Args:
        filename: Name of the file to save to
    
    Returns:
        bool: True if saved successfully
    """
    try:
        state_data = {
            "state": _state_storage,
            "activity_log": _activity_log[-100:],  # Save last 100 activities
            "timestamp": get_timestamp(),
            "version": "1.0"
        }
        
        return save_data(state_data, filename, "json")
        
    except Exception as e:
        logger.error(f"Failed to save state to file: {e}")
        return False

@handle_exceptions
def load_state_from_file(filename: str = "app_state.json") -> bool:
    """
    Load state from file
    
    Args:
        filename: Name of the file to load from
    
    Returns:
        bool: True if loaded successfully
    """
    try:
        global _state_storage, _activity_log
        
        state_data = load_data(filename, "json")
        if state_data:
            _state_storage = safe_get(state_data, "state", {})
            _activity_log = safe_get(state_data, "activity_log", [])
            
            logger.info("State loaded from file successfully")
            record_activity("state_loaded", {"filename": filename})
            return True
        else:
            logger.warning(f"No state file found: {filename}")
            return False
            
    except Exception as e:
        logger.error(f"Failed to load state from file: {e}")
        return False

@handle_exceptions
def get_state_info() -> Dict[str, Any]:
    """
    Get information about current state
    
    Returns:
        Dict: State information
    """
    try:
        return {
            "state_keys": list(_state_storage.keys()),
            "state_size": len(_state_storage),
            "activity_count": len(_activity_log),
            "last_activity": _activity_log[-1] if _activity_log else None,
            "timestamp": get_timestamp()
        }
    except Exception as e:
        logger.error(f"Failed to get state info: {e}")
        return {}

# Session state management for Streamlit compatibility
@handle_exceptions
def init_session_state():
    """
    Initialize session state for Streamlit
    """
    try:
        import streamlit as st
        
        if 'mona_initialized' not in st.session_state:
            st.session_state.mona_initialized = True
            st.session_state.mona_state = {}
            record_activity("session_initialized")
            
    except ImportError:
        # Not in Streamlit environment
        pass
    except Exception as e:
        logger.error(f"Failed to initialize session state: {e}")

@handle_exceptions
def get_session_state(key: str, default: Any = None) -> Any:
    """
    Get value from Streamlit session state
    
    Args:
        key: Session state key
        default: Default value
    
    Returns:
        Any: Session state value or default
    """
    try:
        import streamlit as st
        return st.session_state.get(key, default)
    except ImportError:
        return get_state(key, default)
    except Exception as e:
        logger.error(f"Failed to get session state {key}: {e}")
        return default

@handle_exceptions
def set_session_state(key: str, value: Any) -> bool:
    """
    Set value in Streamlit session state
    
    Args:
        key: Session state key
        value: Value to set
    
    Returns:
        bool: True if set successfully
    """
    try:
        import streamlit as st
        st.session_state[key] = value
        return True
    except ImportError:
        return set_state(key, value)
    except Exception as e:
        logger.error(f"Failed to set session state {key}: {e}")
        return False

# Initialize on module import
record_activity("state_module_loaded")