roa2web-service-auto/reports-app/telegram-bot/app/agent/session.py

"""
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'
]