import json
import time
from typing import Dict, Any, Optional
from core.memory import load_user_state, save_user_state
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class SessionManager:
    """Manages user sessions and conversation context"""
    
    def __init__(self, session_timeout: int = 3600):
        """
        Initialize session manager
        
        Args:
            session_timeout: Session timeout in seconds (default: 1 hour)
        """
        self.session_timeout = session_timeout

    def get_session(self, user_id: str) -> Dict[str, Any]:
        """
        Retrieve user session data
        
        Args:
            user_id: Unique identifier for the user
            
        Returns:
            Dictionary containing session data
        """
        try:
            state = load_user_state(user_id)
            if not state:
                logger.info(f"Creating new session for user {user_id}")
                return self._create_new_session()
            
            # Check if session has expired
            last_activity = float(state.get('last_activity', 0))
            if time.time() - last_activity > self.session_timeout:
                logger.info(f"Session expired for user {user_id}, creating new session")
                return self._create_new_session()
            
            return state
        except Exception as e:
            logger.error(f"Error retrieving session for user {user_id}: {e}")
            return self._create_new_session()

    def update_session(self, user_id: str, data: Dict[str, Any]) -> bool:
        """
        Update user session data
        
        Args:
            user_id: Unique identifier for the user
            data: Data to update in the session
            
        Returns:
            Boolean indicating success
        """
        try:
            # Get existing session
            session = self.get_session(user_id)
            
            # Update with new data
            session.update(data)
            session['last_activity'] = time.time()
            
            # Save updated session
            result = save_user_state(user_id, session)
            if result:
                logger.debug(f"Successfully updated session for user {user_id}")
            else:
                logger.warning(f"Failed to save session for user {user_id}")
            return result
        except Exception as e:
            logger.error(f"Error updating session for user {user_id}: {e}")
            return False

    def clear_session(self, user_id: str) -> bool:
        """
        Clear user session data
        
        Args:
            user_id: Unique identifier for the user
            
        Returns:
            Boolean indicating success
        """
        try:
            result = save_user_state(user_id, {})
            if result:
                logger.info(f"Cleared session for user {user_id}")
            return result
        except Exception as e:
            logger.error(f"Error clearing session for user {user_id}: {e}")
            return False

    def _create_new_session(self) -> Dict[str, Any]:
        """
        Create a new session with default values
        
        Returns:
            Dictionary containing new session data
        """
        session = {
            'conversation': [],
            'preferences': {},
            'last_activity': time.time(),
            'created_at': time.time()
        }
        logger.debug("Created new session")
        return session

# Global session manager instance
session_manager = SessionManager()