tianpu-ems/backend/app/main.py

import logging
import uuid
from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from app.api.router import api_router
from app.api.v1.websocket import start_broadcast_task, stop_broadcast_task
from app.core.config import get_settings
from app.core.cache import get_redis, close_redis
from app.services.simulator import DataSimulator
from app.services.report_scheduler import start_scheduler, stop_scheduler
from app.services.aggregation import start_aggregation_scheduler, stop_aggregation_scheduler
from app.collectors.manager import CollectorManager
from app.collectors.queue import IngestionWorker

settings = get_settings()
customer_config = settings.load_customer_config()
simulator = DataSimulator()
collector_manager: Optional[CollectorManager] = None
ingestion_worker: Optional[IngestionWorker] = None

logger = logging.getLogger("app")


@asynccontextmanager
async def lifespan(app: FastAPI):
    global collector_manager, ingestion_worker

    logger.info("Loading customer: %s (%s)", settings.CUSTOMER,
                 customer_config.get("customer_name", settings.CUSTOMER))

    # Initialize Redis cache
    if settings.REDIS_ENABLED:
        redis = await get_redis()
        if redis:
            logger.info("Redis cache initialized")

    # Start aggregation scheduler
    if settings.AGGREGATION_ENABLED:
        await start_aggregation_scheduler()
        logger.info("Aggregation scheduler started")

    # Start ingestion worker
    if settings.INGESTION_QUEUE_ENABLED:
        ingestion_worker = IngestionWorker()
        await ingestion_worker.start()
        logger.info("Ingestion worker started")

    if settings.USE_SIMULATOR:
        logger.info("Starting in SIMULATOR mode")
        await simulator.start()
    else:
        logger.info("Starting in COLLECTOR mode (real IoT devices)")
        collector_manager = CollectorManager()
        await collector_manager.start()
    start_broadcast_task()
    await start_scheduler()
    yield
    await stop_scheduler()
    stop_broadcast_task()
    if settings.USE_SIMULATOR:
        await simulator.stop()
    else:
        if collector_manager:
            await collector_manager.stop()
            collector_manager = None

    # Stop ingestion worker
    if ingestion_worker:
        await ingestion_worker.stop()
        ingestion_worker = None

    # Stop aggregation scheduler
    if settings.AGGREGATION_ENABLED:
        await stop_aggregation_scheduler()

    # Close Redis
    if settings.REDIS_ENABLED:
        await close_redis()
        logger.info("Redis cache closed")


app = FastAPI(
    title=customer_config.get("platform_name", "天普零碳园区智慧能源管理平台"),
    description=customer_config.get("platform_name_en", "Tianpu Zero-Carbon Park Smart Energy Management System"),
    version="1.0.0",
    lifespan=lifespan,
)

_default_origins = ["http://localhost:3000", "http://localhost:5173", "http://127.0.0.1:3000", "http://127.0.0.1:5173"]
_customer_origins = customer_config.get("cors_origins", [])
_cors_origins = list(set(_default_origins + _customer_origins))

app.add_middleware(
    CORSMiddleware,
    allow_origins=_cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.middleware("http")
async def request_id_middleware(request: Request, call_next):
    """Add a unique X-Request-ID header to every response."""
    request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))
    request.state.request_id = request_id
    response = await call_next(request)
    response.headers["X-Request-ID"] = request_id
    return response


@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    """Global exception handler for consistent error responses."""
    request_id = getattr(request.state, "request_id", "unknown")
    logger.error("Unhandled exception [request_id=%s]: %s", request_id, exc, exc_info=True)
    return JSONResponse(
        status_code=500,
        content={
            "detail": "Internal server error",
            "request_id": request_id,
        },
    )


app.include_router(api_router)


@app.get("/health")
async def health():
    return {"status": "ok", "app": settings.APP_NAME}