main.py

"""
AQI Intelligence Engine — FastAPI Main Application

The central entry point that:
1. Loads all AI models on startup
2. Registers all API routers
3. Serves both locally and on HuggingFace Spaces (Docker)
"""

import logging
import sys
import io

# Fix Windows console encoding for emoji rendering
if sys.platform.startswith("win"):
    try:
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
        sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
    except AttributeError:
        pass

import time
from contextlib import asynccontextmanager
from datetime import datetime, timezone

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from pydantic import BaseModel, Field
from config import SERVER_CONFIG, DEVICE

class HealthResponse(BaseModel):
    status: str = "ok"
    models_loaded: bool = False
    device: str = "cpu"
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    version: str = "2.0.0"

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("aqi_engine")

# Track model loading state
models_loaded = False
startup_time = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load AI models on startup, clean up on shutdown."""
    global models_loaded, startup_time

    logger.info("=" * 60)
    logger.info("🚀 AQI Intelligence Engine — Starting Up")
    logger.info(f"   Device: {DEVICE}")
    logger.info("=" * 60)

    import asyncio
    import anyio

    async def load_models_bg():
        global models_loaded, startup_time
        bg_start = time.time()
        
        try:
            from services.forecast.service import load_model as load_forecast
            await anyio.to_thread.run_sync(load_forecast)
            logger.info("✅ Forecast model loaded successfully")
        except Exception as e:
            logger.warning(f"Forecast model loading deferred: {e}")

        # Heavy vision/segmentation models are loaded on-demand to prevent OOM crashes in container
        logger.info("💤 Vision & Segmentation models deferred to on-demand loading")

        elapsed = round(time.time() - bg_start, 2)
        models_loaded = True
        startup_time = elapsed
        logger.info(f"✅ Startup models loaded in {elapsed}s")

    # Start loading in background without blocking port binding
    asyncio.create_task(load_models_bg())
    
    logger.info(f"📡 API docs: http://localhost:{SERVER_CONFIG['port']}/docs")
    logger.info("=" * 60)

    yield  # Application runs here

    logger.info("Shutting down AQI Intelligence Engine...")


# =============================================================================
# FastAPI Application
# =============================================================================

app = FastAPI(
    title="AQI Intelligence Engine",
    description=(
        "AI-powered Urban Air Quality Intelligence Platform for Smart City Intervention. "
        "Provides AQI forecasting (TimesFM 2.5), satellite vision analysis (Florence-2 + "
        "Grounding DINO), segmentation (SAM2), geospatial analysis, hotspot detection, "
        "atmospheric dispersion modeling, source attribution, health risk assessment, "
        "and inspector route optimization."
    ),
    version="2.0.0",
    lifespan=lifespan,
)

# =============================================================================
# CORS Middleware (allow Next.js frontend to call this API)
# =============================================================================

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # In production: restrict to your Next.js domain
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# =============================================================================
# Register Routers
# =============================================================================

from routers.services import services_router

app.include_router(services_router, prefix="/api/v1")
app.include_router(services_router, prefix="/api")
app.include_router(services_router)


# =============================================================================
# Mount Frontend Static Files & Serve index.html at Root
# =============================================================================

from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import os

# Create frontend directory if it doesn't exist yet
os.makedirs("frontend", exist_ok=True)

# Mount the static directory
app.mount("/frontend", StaticFiles(directory="frontend"), name="frontend")



# =============================================================================
# Health Check
# =============================================================================

@app.get("/health", response_model=HealthResponse, tags=["System"])
async def health_check():
    """Health check endpoint — verifies the server and models are running."""
    return HealthResponse(
        status="ok",
        models_loaded=models_loaded,
        device=DEVICE,
        timestamp=datetime.now(timezone.utc),
        version="2.0.0",
    )


@app.get("/", tags=["System"])
async def root():
    """Root endpoint serving the interactive control center dashboard."""
    index_path = os.path.join("frontend", "index.html")
    if os.path.exists(index_path):
        return FileResponse(index_path)
    return {
        "name": "AQI Intelligence Engine",
        "version": "2.0.0",
        "message": "Frontend dashboard not built yet. Visit /docs for APIs.",
        "docs": "/docs",
        "health": "/health"
    }


# =============================================================================
# Run with: uvicorn main:app --reload --port 7860
# =============================================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "main:app",
        host=SERVER_CONFIG["host"],
        port=SERVER_CONFIG["port"],
        reload=SERVER_CONFIG["reload"],
        workers=SERVER_CONFIG["workers"],
        reload_excludes=["models/*", "**/models/*", "*.json", "frontend/*"],
    )