sac4cps-backend/main.py

import asyncio
import json
import redis.asyncio as redis
import time
import os
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Depends, Query
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Optional
import logging
from contextlib import asynccontextmanager

# Import our custom modules
from database import connect_to_mongo, close_mongo_connection, redis_manager, schedule_cleanup
from persistence import persistence_service
from models import DataQuery, DataResponse, HealthCheck
from api import router as api_router

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Application startup time for uptime calculation
app_start_time = time.time()

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager"""
    # Startup
    logger.info("Application starting up...")
    
    # Connect to databases
    await connect_to_mongo()
    await persistence_service.initialize()
    
    # Start background tasks
    asyncio.create_task(redis_subscriber())
    asyncio.create_task(schedule_cleanup())
    
    logger.info("Application startup complete")
    
    yield
    
    # Shutdown
    logger.info("Application shutting down...")
    await close_mongo_connection()
    await redis_manager.disconnect()
    logger.info("Application shutdown complete")

app = FastAPI(
    title="Energy Monitoring Dashboard API",
    description="Real-time energy monitoring and IoT sensor data management system",
    version="1.0.0",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Configure appropriately for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include API router
app.include_router(api_router, prefix="/api/v1")

# In-memory store for active WebSocket connections
active_connections: List[WebSocket] = []

# Redis channel to subscribe to
REDIS_CHANNEL = "energy_data"


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """
    WebSocket endpoint that connects a client, adds them to the active pool,
    and removes them on disconnection.
    """
    await websocket.accept()
    active_connections.append(websocket)
    logger.info(f"New client connected. Total clients: {len(active_connections)}")
    try:
        while True:
            # Keep the connection alive
            await websocket.receive_text()
    except WebSocketDisconnect:
        active_connections.remove(websocket)
        logger.info(f"Client disconnected. Total clients: {len(active_connections)}")


async def redis_subscriber():
    """
    Connects to Redis, subscribes to the specified channel, and broadcasts
    messages to all active WebSocket clients. Also persists data to MongoDB.
    """
    logger.info("Starting Redis subscriber...")
    try:
        r = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)
        await r.ping()
        logger.info("Successfully connected to Redis for subscription.")
    except Exception as e:
        logger.error(f"Could not connect to Redis for subscription: {e}")
        return

    pubsub = r.pubsub()
    await pubsub.subscribe(REDIS_CHANNEL)
    
    logger.info(f"Subscribed to Redis channel: '{REDIS_CHANNEL}'")
    while True:
        try:
            message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
            if message:
                message_data = message['data']
                logger.debug(f"Received from Redis: {message_data}")
                
                # Process and persist the data
                await persistence_service.process_sensor_message(message_data)
                
                # Broadcast message to all connected WebSocket clients
                if active_connections:
                    await asyncio.gather(
                        *[connection.send_text(message_data) for connection in active_connections],
                        return_exceptions=True
                    )
                
        except Exception as e:
            logger.error(f"Error in Redis subscriber loop: {e}")
            # Add a delay to prevent rapid-fire errors
            await asyncio.sleep(5)


@app.get("/")
async def read_root():
    """Root endpoint with basic system information"""
    return {
        "message": "Energy Monitoring Dashboard Backend",
        "version": "1.0.0",
        "status": "running",
        "uptime_seconds": time.time() - app_start_time
    }


@app.get("/health", response_model=HealthCheck)
async def health_check():
    """Health check endpoint"""
    try:
        # Check database connections
        mongodb_connected = True
        redis_connected = True
        
        try:
            await persistence_service.db.command("ping")
        except:
            mongodb_connected = False
            
        try:
            await redis_manager.redis_client.ping()
        except:
            redis_connected = False
        
        # Get system statistics
        stats = await persistence_service.get_sensor_statistics()
        
        # Determine overall status
        status = "healthy"
        if not mongodb_connected or not redis_connected:
            status = "degraded"
        
        return HealthCheck(
            status=status,
            mongodb_connected=mongodb_connected,
            redis_connected=redis_connected,
            total_sensors=stats.get("total_sensors", 0),
            active_sensors=stats.get("active_sensors", 0),
            total_readings=stats.get("total_readings", 0),
            uptime_seconds=time.time() - app_start_time
        )
        
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        raise HTTPException(status_code=503, detail="Service Unavailable")


@app.get("/status")
async def system_status():
    """Detailed system status endpoint"""
    try:
        stats = await persistence_service.get_sensor_statistics()
        
        return {
            "timestamp": time.time(),
            "uptime_seconds": time.time() - app_start_time,
            "active_websocket_connections": len(active_connections),
            "database_stats": stats
        }
        
    except Exception as e:
        logger.error(f"Status check failed: {e}")
        raise HTTPException(status_code=500, detail="Internal Server Error")