python-detector-worker/core/communication/session_integration.py

"""
Integration layer between WebSocket handler and Session Process Manager.
Bridges the existing WebSocket protocol with the new session-based architecture.
"""

import asyncio
import logging
from typing import Dict, Any, Optional
import numpy as np

from ..processes.session_manager import SessionProcessManager
from ..processes.communication import DetectionResultResponse, ErrorResponse
from .state import worker_state
from .messages import serialize_outgoing_message
# Streaming is now handled directly by session workers - no shared stream manager needed

logger = logging.getLogger(__name__)


class SessionWebSocketIntegration:
    """
    Integration layer that connects WebSocket protocol with Session Process Manager.
    Maintains compatibility with existing WebSocket message handling.
    """

    def __init__(self, websocket_handler=None):
        """
        Initialize session WebSocket integration.

        Args:
            websocket_handler: Reference to WebSocket handler for sending messages
        """
        self.websocket_handler = websocket_handler
        self.session_manager = SessionProcessManager()

        # Track active subscriptions for compatibility
        self.active_subscriptions: Dict[str, Dict[str, Any]] = {}

        # Set up callbacks
        self.session_manager.set_detection_result_callback(self._on_detection_result)
        self.session_manager.set_error_callback(self._on_session_error)

    async def start(self):
        """Start the session integration."""
        await self.session_manager.start()
        logger.info("Session WebSocket integration started")

    async def stop(self):
        """Stop the session integration."""
        await self.session_manager.stop()
        logger.info("Session WebSocket integration stopped")

    async def handle_set_subscription_list(self, message) -> bool:
        """
        Handle setSubscriptionList message by managing session processes.

        Args:
            message: SetSubscriptionListMessage

        Returns:
            True if successful
        """
        try:
            logger.info(f"Processing subscription list with {len(message.subscriptions)} subscriptions")

            new_subscription_ids = set()
            for subscription in message.subscriptions:
                subscription_id = subscription.subscriptionIdentifier
                new_subscription_ids.add(subscription_id)

                # Check if this is a new subscription
                if subscription_id not in self.active_subscriptions:
                    logger.info(f"Creating new session for subscription: {subscription_id}")

                    # Convert subscription to configuration dict
                    subscription_config = {
                        'subscriptionIdentifier': subscription.subscriptionIdentifier,
                        'rtspUrl': getattr(subscription, 'rtspUrl', None),
                        'snapshotUrl': getattr(subscription, 'snapshotUrl', None),
                        'snapshotInterval': getattr(subscription, 'snapshotInterval', 5000),
                        'modelUrl': subscription.modelUrl,
                        'modelId': subscription.modelId,
                        'modelName': subscription.modelName,
                        'cropX1': subscription.cropX1,
                        'cropY1': subscription.cropY1,
                        'cropX2': subscription.cropX2,
                        'cropY2': subscription.cropY2
                    }

                    # Create session process
                    success = await self.session_manager.create_session(
                        subscription_id, subscription_config
                    )

                    if success:
                        self.active_subscriptions[subscription_id] = subscription_config
                        logger.info(f"Session created successfully for {subscription_id}")

                        # Stream handling is now integrated into session worker process
                    else:
                        logger.error(f"Failed to create session for {subscription_id}")
                        return False

                else:
                    # Update existing subscription configuration if needed
                    self.active_subscriptions[subscription_id].update({
                        'modelUrl': subscription.modelUrl,
                        'modelId': subscription.modelId,
                        'modelName': subscription.modelName,
                        'cropX1': subscription.cropX1,
                        'cropY1': subscription.cropY1,
                        'cropX2': subscription.cropX2,
                        'cropY2': subscription.cropY2
                    })

            # Remove sessions for subscriptions that are no longer active
            current_subscription_ids = set(self.active_subscriptions.keys())
            removed_subscriptions = current_subscription_ids - new_subscription_ids

            for subscription_id in removed_subscriptions:
                logger.info(f"Removing session for subscription: {subscription_id}")
                await self.session_manager.remove_session(subscription_id)
                del self.active_subscriptions[subscription_id]

            # Update worker state for compatibility
            worker_state.set_subscriptions(message.subscriptions)

            logger.info(f"Subscription list processed: {len(new_subscription_ids)} active sessions")
            return True

        except Exception as e:
            logger.error(f"Error handling subscription list: {e}", exc_info=True)
            return False

    async def handle_set_session_id(self, message) -> bool:
        """
        Handle setSessionId message by forwarding to appropriate session process.

        Args:
            message: SetSessionIdMessage

        Returns:
            True if successful
        """
        try:
            display_id = message.payload.displayIdentifier
            session_id = message.payload.sessionId

            logger.info(f"Setting session ID {session_id} for display {display_id}")

            # Find subscription identifier for this display
            subscription_id = None
            for sub_id in self.active_subscriptions.keys():
                # Extract display identifier from subscription identifier
                if display_id in sub_id:
                    subscription_id = sub_id
                    break

            if not subscription_id:
                logger.error(f"No active subscription found for display {display_id}")
                return False

            # Forward to session process
            success = await self.session_manager.set_session_id(
                subscription_id, str(session_id), display_id
            )

            if success:
                # Update worker state for compatibility
                worker_state.set_session_id(display_id, session_id)
                logger.info(f"Session ID {session_id} set successfully for {display_id}")
            else:
                logger.error(f"Failed to set session ID {session_id} for {display_id}")

            return success

        except Exception as e:
            logger.error(f"Error setting session ID: {e}", exc_info=True)
            return False

    async def process_frame(self, subscription_id: str, frame: np.ndarray, display_id: str, timestamp: float = None) -> bool:
        """
        Process frame through appropriate session process.

        Args:
            subscription_id: Subscription identifier
            frame: Frame to process
            display_id: Display identifier
            timestamp: Frame timestamp

        Returns:
            True if frame was processed successfully
        """
        try:
            if timestamp is None:
                timestamp = asyncio.get_event_loop().time()

            # Forward frame to session process
            success = await self.session_manager.process_frame(
                subscription_id, frame, display_id, timestamp
            )

            if not success:
                logger.warning(f"Failed to process frame for subscription {subscription_id}")

            return success

        except Exception as e:
            logger.error(f"Error processing frame for {subscription_id}: {e}", exc_info=True)
            return False

    async def _on_detection_result(self, subscription_id: str, response: DetectionResultResponse):
        """
        Handle detection result from session process.

        Args:
            subscription_id: Subscription identifier
            response: Detection result response
        """
        try:
            logger.debug(f"Received detection result from {subscription_id}: phase={response.phase}")

            # Send imageDetection message via WebSocket (if needed)
            if self.websocket_handler and hasattr(self.websocket_handler, 'send_message'):
                from .models import ImageDetectionMessage, DetectionData

                # Convert response detections to the expected format
                # The DetectionData expects modelId and modelName, and detection dict
                detection_data = DetectionData(
                    detection=response.detections,
                    modelId=getattr(response, 'model_id', 0),  # Get from response if available
                    modelName=getattr(response, 'model_name', 'unknown')  # Get from response if available
                )

                # Convert timestamp to string format if it exists
                timestamp_str = None
                if hasattr(response, 'timestamp') and response.timestamp:
                    from datetime import datetime
                    if isinstance(response.timestamp, (int, float)):
                        # Convert Unix timestamp to ISO format string
                        timestamp_str = datetime.fromtimestamp(response.timestamp).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
                    else:
                        timestamp_str = str(response.timestamp)

                detection_message = ImageDetectionMessage(
                    subscriptionIdentifier=subscription_id,
                    data=detection_data,
                    timestamp=timestamp_str
                )

                serialized = serialize_outgoing_message(detection_message)
                await self.websocket_handler.send_message(serialized)

        except Exception as e:
            logger.error(f"Error handling detection result from {subscription_id}: {e}", exc_info=True)

    async def _on_session_error(self, subscription_id: str, error_response: ErrorResponse):
        """
        Handle error from session process.

        Args:
            subscription_id: Subscription identifier
            error_response: Error response
        """
        logger.error(f"Session error from {subscription_id}: {error_response.error_type} - {error_response.error_message}")

        # Send error message via WebSocket if needed
        if self.websocket_handler and hasattr(self.websocket_handler, 'send_message'):
            error_message = {
                'type': 'sessionError',
                'payload': {
                    'subscriptionIdentifier': subscription_id,
                    'errorType': error_response.error_type,
                    'errorMessage': error_response.error_message,
                    'timestamp': error_response.timestamp
                }
            }

            try:
                serialized = serialize_outgoing_message(error_message)
                await self.websocket_handler.send_message(serialized)
            except Exception as e:
                logger.error(f"Failed to send error message: {e}")

    def get_session_stats(self) -> Dict[str, Any]:
        """
        Get statistics about active sessions.

        Returns:
            Dictionary with session statistics
        """
        return {
            'active_sessions': self.session_manager.get_session_count(),
            'max_sessions': self.session_manager.max_concurrent_sessions,
            'subscriptions': list(self.active_subscriptions.keys())
        }

    async def handle_progression_stage(self, message) -> bool:
        """
        Handle setProgressionStage message.

        Args:
            message: SetProgressionStageMessage

        Returns:
            True if successful
        """
        try:
            # For now, just update worker state for compatibility
            # In future phases, this could be forwarded to session processes
            worker_state.set_progression_stage(
                message.payload.displayIdentifier,
                message.payload.progressionStage
            )
            return True
        except Exception as e:
            logger.error(f"Error handling progression stage: {e}", exc_info=True)
            return False