python-detector-worker/app.py

"""
Detector Worker - Main FastAPI Application
Refactored modular architecture for computer vision pipeline processing.
"""
import json
import logging
import os
import time
import cv2
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, HTTPException
from fastapi.responses import Response

# Import new modular communication system
from core.communication.websocket import websocket_endpoint
from core.communication.state import worker_state

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=[
        logging.FileHandler("detector_worker.log"),
        logging.StreamHandler()
    ]
)

logger = logging.getLogger("detector_worker")
logger.setLevel(logging.DEBUG)

# Frames are now stored in the shared cache buffer from core.streaming.buffers
# latest_frames = {}  # Deprecated - using shared_cache_buffer instead

# Lifespan event handler (modern FastAPI approach)
@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan management."""
    # Startup
    logger.info("Detector Worker started successfully")
    logger.info("WebSocket endpoint available at: ws://0.0.0.0:8001/")
    logger.info("HTTP camera endpoint available at: http://0.0.0.0:8001/camera/{camera_id}/image")
    logger.info("Health check available at: http://0.0.0.0:8001/health")
    logger.info("Ready and waiting for backend WebSocket connections")

    yield

    # Shutdown
    logger.info("Detector Worker shutting down...")
    # Clear all state
    worker_state.set_subscriptions([])
    worker_state.session_ids.clear()
    worker_state.progression_stages.clear()
    # latest_frames.clear()  # No longer needed - frames are in shared_cache_buffer
    logger.info("Detector Worker shutdown complete")

# Create FastAPI application with detailed WebSocket logging
app = FastAPI(title="Detector Worker", version="2.0.0", lifespan=lifespan)

# Add middleware to log all requests
@app.middleware("http")
async def log_requests(request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    logger.debug(f"HTTP {request.method} {request.url} - {response.status_code} ({process_time:.3f}s)")
    return response

# Load configuration
config_path = "config.json"
if os.path.exists(config_path):
    with open(config_path, "r") as f:
        config = json.load(f)
    logger.info(f"Loaded configuration from {config_path}")
else:
    # Default configuration
    config = {
        "poll_interval_ms": 100,
        "reconnect_interval_sec": 5,
        "target_fps": 10,
        "max_streams": 20,
        "max_retries": 3
    }
    logger.warning(f"Configuration file {config_path} not found, using defaults")

# Ensure models directory exists
os.makedirs("models", exist_ok=True)
logger.info("Ensured models directory exists")

# Initialize stream manager with config value
from core.streaming import initialize_stream_manager
initialize_stream_manager(max_streams=config.get('max_streams', 10))
logger.info(f"Initialized stream manager with max_streams={config.get('max_streams', 10)}")

# Frames are now stored in the shared cache buffer from core.streaming.buffers
# latest_frames = {}  # Deprecated - using shared_cache_buffer instead

logger.info("Starting detector worker application (refactored)")
logger.info(f"Configuration: Target FPS: {config.get('target_fps', 10)}, "
           f"Max streams: {config.get('max_streams', 5)}, "
           f"Max retries: {config.get('max_retries', 3)}")


@app.websocket("/")
async def websocket_handler(websocket: WebSocket):
    """
    Main WebSocket endpoint for backend communication.
    Handles all protocol messages according to worker.md specification.
    """
    client_info = f"{websocket.client.host}:{websocket.client.port}" if websocket.client else "unknown"
    logger.info(f"[RX ← Backend] New WebSocket connection request from {client_info}")

    try:
        await websocket_endpoint(websocket)
    except Exception as e:
        logger.error(f"WebSocket handler error for {client_info}: {e}", exc_info=True)


@app.get("/camera/{camera_id}/image")
async def get_camera_image(camera_id: str):
    """
    HTTP endpoint to retrieve the latest frame from a camera as JPEG image.

    This endpoint is preserved for backward compatibility with existing systems.

    Args:
        camera_id: The subscription identifier (e.g., "display-001;cam-001")

    Returns:
        JPEG image as binary response

    Raises:
        HTTPException: 404 if camera not found or no frame available
        HTTPException: 500 if encoding fails
    """
    try:
        from urllib.parse import unquote

        # URL decode the camera_id to handle encoded characters
        original_camera_id = camera_id
        camera_id = unquote(camera_id)
        logger.debug(f"REST API request: original='{original_camera_id}', decoded='{camera_id}'")

        # Check if camera is in active subscriptions
        subscription = worker_state.get_subscription(camera_id)
        if not subscription:
            logger.warning(f"Camera ID '{camera_id}' not found in active subscriptions")
            available_cameras = list(worker_state.subscriptions.keys())
            logger.debug(f"Available cameras: {available_cameras}")
            raise HTTPException(
                status_code=404,
                detail=f"Camera {camera_id} not found or not active"
            )

        # Extract actual camera_id from subscription identifier (displayId;cameraId)
        # Frames are stored using just the camera_id part
        actual_camera_id = camera_id.split(';')[-1] if ';' in camera_id else camera_id

        # Get frame from the shared cache buffer
        from core.streaming.buffers import shared_cache_buffer

        # Debug: Log available cameras in buffer
        available_cameras = shared_cache_buffer.frame_buffer.get_camera_list()
        logger.debug(f"Available cameras in buffer: {available_cameras}")
        logger.debug(f"Looking for camera: '{actual_camera_id}'")

        frame = shared_cache_buffer.get_frame(actual_camera_id)
        if frame is None:
            logger.warning(f"No cached frame available for camera '{actual_camera_id}' (from subscription '{camera_id}')")
            logger.warning(f"Available cameras in buffer: {available_cameras}")
            raise HTTPException(
                status_code=404,
                detail=f"No frame available for camera {actual_camera_id}"
            )

        logger.debug(f"Retrieved cached frame for camera '{actual_camera_id}' (from subscription '{camera_id}'), shape: {frame.shape}")

        # Encode frame as JPEG
        success, buffer_img = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
        if not success:
            raise HTTPException(status_code=500, detail="Failed to encode image as JPEG")

        # Return image as binary response
        return Response(content=buffer_img.tobytes(), media_type="image/jpeg")

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error retrieving image for camera {camera_id}: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")


@app.get("/health")
async def health_check():
    """Health check endpoint for monitoring."""
    return {
        "status": "healthy",
        "version": "2.0.0",
        "active_subscriptions": len(worker_state.subscriptions),
        "active_sessions": len(worker_state.session_ids)
    }




if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001)