api/main.py

"""
api/main.py
FastAPI application — mounts all routers, configures CORS, rate limiting, lifespan.
"""

import os
from contextlib import asynccontextmanager
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded

from api.routers import query, upload, history, dashboard, report, schema
from api.routers import metrics as metrics_router
from api.routers import profile as profile_router


# ── Rate limiter ──────────────────────────────────────────────────────────────
limiter = Limiter(key_func=get_remote_address)

import asyncio
import httpx

async def keep_awake():
    """Background task to ping the server and prevent HF Space from sleeping."""
    while True:
        try:
            await asyncio.sleep(600)  # Ping every 10 minutes
            async with httpx.AsyncClient() as client:
                await client.get("http://127.0.0.1:7860/health", timeout=10.0)
                print("[Keep-Awake] Pinged health endpoint to prevent sleep.")
        except asyncio.CancelledError:
            break
        except Exception as e:
            print(f"[Keep-Awake] Ping failed: {e}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: pre-warm LLM clients + embedding model
    from llm import get_groq_client, get_embedder
    get_groq_client()
    # Pre-download and load embedding model (first run downloads ~420MB)
    print("[Lifespan] Loading embedding model (first run may download ~420MB)...")
    get_embedder()
    print("[Lifespan] Embedding model loaded successfully.")
    
    # Run automatic database migrations on boot
    print("[Lifespan] Ensuring database tables exist...")
    try:
        from db.pool import get_connection, release_connection
        import os
        conn = get_connection()
        try:
            with conn.cursor() as cur:
                with open(os.path.join(os.path.dirname(os.path.dirname(__file__)), "scripts", "migrate.sql"), "r") as f:
                    cur.execute(f.read())
                conn.commit()
            print("[Lifespan] Database migrations executed successfully.")
        finally:
            release_connection(conn)
    except Exception as e:
        print(f"[Lifespan] Failed to run database migrations: {e}")
    
    # Start the keep-awake background task
    ping_task = asyncio.create_task(keep_awake())
    
    yield
    
    # Shutdown: cancel task and close DB pool
    ping_task.cancel()
    try:
        from db.pool import get_pool
        get_pool().closeall()
    except Exception:
        pass


app = FastAPI(
    title="Cloud Data Analyst Agent",
    version="2.0.0",
    description="AI-powered data analyst with self-correcting LangGraph agent, "
                "multi-turn conversations, anomaly detection, and real-time observability.",
    docs_url="/docs",
    lifespan=lifespan,
)

# Attach rate limiter
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

# CORS — allow the Vite frontend and any Render preview URL
ALLOWED_ORIGINS = os.environ.get(
    "ALLOWED_ORIGINS",
    "http://localhost:5173,http://localhost:3000",
).split(",")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"] if os.environ.get("DEMO_MODE") == "true" else ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Routers
app.include_router(query.router,            prefix="/api/query",     tags=["query"])
app.include_router(upload.router,           prefix="/api/upload",    tags=["upload"])
app.include_router(history.router,          prefix="/api/history",   tags=["history"])
app.include_router(dashboard.router,        prefix="/api/dashboard", tags=["dashboard"])
app.include_router(report.router,           prefix="/api/report",    tags=["report"])
app.include_router(schema.router,           prefix="/api/schema",    tags=["schema"])
app.include_router(metrics_router.router,   prefix="/api/metrics",   tags=["metrics"])
app.include_router(profile_router.router,   prefix="/api/profile",   tags=["profile"])


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

# ── Serve Frontend ────────────────────────────────────────────────────────────
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

frontend_dist = os.path.join(os.path.dirname(os.path.dirname(__file__)), "frontend", "dist")

# Mount the assets directory specifically if it exists
assets_dir = os.path.join(frontend_dist, "assets")
if os.path.exists(assets_dir):
    app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")
    
@app.get("/{full_path:path}")
async def serve_frontend(full_path: str):
    # Prevent API routes from being intercepted if they fall through
    if full_path.startswith("api/"):
        from fastapi import HTTPException
        raise HTTPException(status_code=404, detail="API route not found")
        
    file_path = os.path.join(frontend_dist, full_path)
    if full_path and os.path.isfile(file_path):
        return FileResponse(file_path)
    # Fallback to index.html for React Router
    index_file = os.path.join(frontend_dist, "index.html")
    if os.path.isfile(index_file):
        return FileResponse(index_file)
    return {"error": "Frontend build not found"}