""" Session Management for Telegram Bot This module handles session state for Telegram users, specifically managing the active company selection for command handlers. """ import logging import json from typing import Dict, Any, Optional from datetime import datetime from app.db.operations import ( create_session, get_user_active_session, update_session_state, delete_user_sessions ) logger = logging.getLogger(__name__) class ConversationSession: """ Manages session state for a single user. Attributes: telegram_user_id: Telegram user ID session_id: UUID of the session active_company_id: Selected company ID active_company_name: Selected company name active_company_cui: Selected company CUI """ def __init__( self, telegram_user_id: int, session_id: Optional[str] = None ): """ Initialize a session. Args: telegram_user_id: Telegram user ID session_id: Existing session ID (if resuming), or None for new session """ self.telegram_user_id = telegram_user_id self.session_id = session_id self.created_at = datetime.now() self.updated_at = datetime.now() # Active company for this session self.active_company_id: Optional[int] = None self.active_company_name: Optional[str] = None self.active_company_cui: Optional[str] = None def set_active_company( self, company_id: int, company_name: str, company_cui: Optional[str] = None ): """ Set the active company for this session. Args: company_id: Company ID company_name: Company name company_cui: Company CUI (optional) """ self.active_company_id = company_id self.active_company_name = company_name self.active_company_cui = company_cui self.updated_at = datetime.now() logger.info( f"Active company set for user {self.telegram_user_id}: " f"{company_name} (ID: {company_id})" ) def get_active_company(self) -> Optional[Dict[str, Any]]: """ Get the active company information. Returns: Dict with company info (id, name, cui) or None if no company selected """ if self.active_company_id is not None: return { "id": self.active_company_id, "name": self.active_company_name, "cui": self.active_company_cui } return None def clear_active_company(self): """ Clear the active company selection. """ logger.info( f"Clearing active company for user {self.telegram_user_id} " f"(was: {self.active_company_name})" ) self.active_company_id = None self.active_company_name = None self.active_company_cui = None self.updated_at = datetime.now() def to_dict(self) -> Dict[str, Any]: """ Serialize session to dictionary (for database storage). Returns: Dict representation of session """ return { "telegram_user_id": self.telegram_user_id, "session_id": self.session_id, "created_at": self.created_at.isoformat(), "updated_at": self.updated_at.isoformat(), "active_company_id": self.active_company_id, "active_company_name": self.active_company_name, "active_company_cui": self.active_company_cui } @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'ConversationSession': """ Deserialize session from dictionary. Args: data: Dict representation of session Returns: ConversationSession instance """ session = cls( telegram_user_id=data["telegram_user_id"], session_id=data.get("session_id") ) # Restore active company session.active_company_id = data.get("active_company_id") session.active_company_name = data.get("active_company_name") session.active_company_cui = data.get("active_company_cui") if "created_at" in data: session.created_at = datetime.fromisoformat(data["created_at"]) if "updated_at" in data: session.updated_at = datetime.fromisoformat(data["updated_at"]) return session class SessionManager: """ Manages sessions for all users. Provides methods to create, retrieve, update, and delete sessions. Sessions are stored both in memory (for quick access) and in database (for persistence). """ def __init__(self): """ Initialize the session manager. """ self._sessions: Dict[int, ConversationSession] = {} logger.info("SessionManager initialized") async def get_or_create_session( self, telegram_user_id: int ) -> ConversationSession: """ Get existing session for a user or create a new one. Args: telegram_user_id: Telegram user ID Returns: ConversationSession for the user """ # Check in-memory cache first if telegram_user_id in self._sessions: logger.debug(f"Found session in cache for user {telegram_user_id}") return self._sessions[telegram_user_id] # Check database for existing session session_data = await get_user_active_session(telegram_user_id) if session_data: # Restore session from database conversation_state_json = session_data.get('conversation_state') if conversation_state_json: try: session_dict = json.loads(conversation_state_json) session = ConversationSession.from_dict(session_dict) session.session_id = session_data['session_id'] self._sessions[telegram_user_id] = session logger.info(f"Restored session from database for user {telegram_user_id}") return session except json.JSONDecodeError as e: logger.error(f"Failed to parse session state: {e}") # Create new session session = ConversationSession(telegram_user_id) # Save to database session_id = await create_session( telegram_user_id=telegram_user_id, conversation_state=json.dumps(session.to_dict()), expires_in_hours=24 ) session.session_id = session_id self._sessions[telegram_user_id] = session logger.info(f"Created new session for user {telegram_user_id} (ID: {session_id})") return session async def save_session(self, telegram_user_id: int) -> bool: """ Save session to database. Args: telegram_user_id: Telegram user ID Returns: bool: True if saved successfully """ session = self._sessions.get(telegram_user_id) if not session or not session.session_id: logger.warning(f"No session to save for user {telegram_user_id}") return False try: conversation_state = json.dumps(session.to_dict()) success = await update_session_state( session_id=session.session_id, conversation_state=conversation_state ) if success: logger.debug(f"Saved session for user {telegram_user_id}") else: logger.warning(f"Failed to save session for user {telegram_user_id}") return success except Exception as e: logger.error(f"Error saving session for user {telegram_user_id}: {e}") return False async def delete_session(self, telegram_user_id: int) -> bool: """ Delete session completely (from memory and database). Args: telegram_user_id: Telegram user ID Returns: bool: True if deleted successfully """ # Remove from memory if telegram_user_id in self._sessions: del self._sessions[telegram_user_id] # Delete from database success = await delete_user_sessions(telegram_user_id) if success: logger.info(f"Deleted session for user {telegram_user_id}") else: logger.warning(f"Failed to delete session for user {telegram_user_id}") return success def get_active_sessions_count(self) -> int: """ Get count of active sessions in memory. Returns: int: Number of active sessions """ return len(self._sessions) # Singleton instance _session_manager_instance: Optional[SessionManager] = None def get_session_manager() -> SessionManager: """ Get or create the singleton SessionManager instance. Returns: SessionManager: Singleton instance """ global _session_manager_instance if _session_manager_instance is None: _session_manager_instance = SessionManager() return _session_manager_instance # Export main classes and functions __all__ = [ 'ConversationSession', 'SessionManager', 'get_session_manager' ]