bakery-ia/services/notification/app/main.py

# ================================================================
# services/notification/app/main.py - ENHANCED WITH SSE SUPPORT
# ================================================================
"""
Notification Service Main Application
Handles email, WhatsApp notifications and SSE for real-time alerts/recommendations
"""

import structlog
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from app.core.config import settings
from app.core.database import init_db
from app.api.notifications import router as notification_router
from app.api.sse_routes import router as sse_router
from app.services.messaging import setup_messaging, cleanup_messaging
from app.services.sse_service import SSEService
from app.services.notification_orchestrator import NotificationOrchestrator
from app.services.email_service import EmailService
from app.services.whatsapp_service import WhatsAppService
from shared.monitoring import setup_logging, HealthChecker
from shared.monitoring.metrics import setup_metrics_early

# Setup logging first
setup_logging("notification-service", settings.LOG_LEVEL)
logger = structlog.get_logger()

# Global variables for lifespan access
metrics_collector = None
health_checker = None

# Create FastAPI app FIRST
app = FastAPI(
    title="Bakery Notification Service",
    description="Email, WhatsApp and SSE notification service for bakery alerts and recommendations",
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# Setup metrics BEFORE any middleware and BEFORE lifespan
metrics_collector = setup_metrics_early(app, "notification-service")

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan events - NO MIDDLEWARE ADDED HERE"""
    global health_checker
    
    # Startup
    logger.info("Starting Notification Service...")
    
    try:
        # Initialize database
        await init_db()
        logger.info("Database initialized")
        
        # Setup messaging
        await setup_messaging()
        logger.info("Messaging initialized")
        
        # Initialize services
        email_service = EmailService()
        whatsapp_service = WhatsAppService()
        
        # Initialize SSE service
        sse_service = SSEService(settings.REDIS_URL)
        await sse_service.initialize()
        logger.info("SSE service initialized")
        
        # Create orchestrator
        orchestrator = NotificationOrchestrator(
            email_service=email_service,
            whatsapp_service=whatsapp_service,
            sse_service=sse_service
        )
        
        # Store services in app state
        app.state.orchestrator = orchestrator
        app.state.sse_service = sse_service
        app.state.email_service = email_service
        app.state.whatsapp_service = whatsapp_service
        
        # Register custom metrics (metrics_collector already exists)
        metrics_collector.register_counter("notifications_sent_total", "Total notifications sent", labels=["type", "status", "channel"])
        metrics_collector.register_counter("emails_sent_total", "Total emails sent", labels=["status"])
        metrics_collector.register_counter("whatsapp_sent_total", "Total WhatsApp messages sent", labels=["status"])
        metrics_collector.register_counter("sse_events_sent_total", "Total SSE events sent", labels=["tenant", "event_type"])
        metrics_collector.register_histogram("notification_processing_duration_seconds", "Time spent processing notifications")
        metrics_collector.register_gauge("notification_queue_size", "Current notification queue size")
        metrics_collector.register_gauge("sse_active_connections", "Number of active SSE connections")
        
        # Setup health checker
        health_checker = HealthChecker("notification-service")
        
        # Add database health check
        async def check_database():
            try:
                from app.core.database import get_db
                from sqlalchemy import text
                async for db in get_db():
                    await db.execute(text("SELECT 1"))
                    return True
            except Exception as e:
                return f"Database error: {e}"
        
        health_checker.add_check("database", check_database, timeout=5.0, critical=True)
        
        # Add email service health check
        async def check_email_service():
            try:
                from app.services.email_service import EmailService
                email_service = EmailService()
                return await email_service.health_check()
            except Exception as e:
                return f"Email service error: {e}"
        
        health_checker.add_check("email_service", check_email_service, timeout=10.0, critical=True)
        
        # Add WhatsApp service health check
        async def check_whatsapp_service():
            try:
                return await whatsapp_service.health_check()
            except Exception as e:
                return f"WhatsApp service error: {e}"
        
        health_checker.add_check("whatsapp_service", check_whatsapp_service, timeout=10.0, critical=False)
        
        # Add SSE service health check
        async def check_sse_service():
            try:
                metrics = sse_service.get_metrics()
                return "healthy" if metrics["redis_connected"] else "Redis connection failed"
            except Exception as e:
                return f"SSE service error: {e}"
        
        health_checker.add_check("sse_service", check_sse_service, timeout=5.0, critical=True)
        
        # Add messaging health check
        def check_messaging():
            try:
                # Check if messaging is properly initialized
                from app.services.messaging import notification_publisher
                return notification_publisher.connected if notification_publisher else False
            except Exception as e:
                return f"Messaging error: {e}"
        
        health_checker.add_check("messaging", check_messaging, timeout=3.0, critical=False)
        
        # Store health checker in app state
        app.state.health_checker = health_checker
        
        logger.info("Notification Service with SSE support started successfully")
        
    except Exception as e:
        logger.error(f"Failed to start Notification Service: {e}")
        raise
    
    yield
    
    # Shutdown
    logger.info("Shutting down Notification Service...")
    try:
        # Shutdown SSE service
        if hasattr(app.state, 'sse_service'):
            await app.state.sse_service.shutdown()
            logger.info("SSE service shutdown completed")
            
        await cleanup_messaging()
        logger.info("Messaging cleanup completed")
    except Exception as e:
        logger.error(f"Error during shutdown: {e}")

# Set lifespan AFTER metrics setup
app.router.lifespan_context = lifespan

# CORS middleware (added after metrics setup)
app.add_middleware(
    CORSMiddleware,
    allow_origins=getattr(settings, 'CORS_ORIGINS', ["*"]),
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include routers
app.include_router(notification_router, prefix="/api/v1", tags=["notifications"])
app.include_router(sse_router, prefix="/api/v1", tags=["sse"])

# Health check endpoint
@app.get("/health")
async def health_check():
    """Comprehensive health check endpoint including SSE"""
    if health_checker:
        health_result = await health_checker.check_health()
        
        # Add SSE metrics to health check
        if hasattr(app.state, 'sse_service'):
            try:
                sse_metrics = app.state.sse_service.get_metrics()
                health_result['sse_metrics'] = sse_metrics
            except Exception as e:
                health_result['sse_error'] = str(e)
        
        return health_result
    else:
        return {
            "service": "notification-service",
            "status": "healthy",
            "version": "2.0.0",
            "features": ["email", "whatsapp", "sse", "alerts", "recommendations"]
        }

# Metrics endpoint
@app.get("/metrics")
async def metrics():
    """Prometheus metrics endpoint"""
    if metrics_collector:
        return metrics_collector.get_metrics()
    return {"metrics": "not_available"}

# Exception handlers
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    """Global exception handler with metrics"""
    logger.error(f"Unhandled exception: {exc}", exc_info=True)
    
    # Record error metric if available
    if metrics_collector:
        metrics_collector.increment_counter("errors_total", labels={"type": "unhandled"})
    
    return JSONResponse(
        status_code=500,
        content={"detail": "Internal server error"}
    )

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