behavior_backend/main.py

import uvicorn
from fastapi import FastAPI, APIRouter, Request, Depends, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import asyncio
from concurrent.futures import ThreadPoolExecutor
from fastapi.responses import JSONResponse
import logging
import time
import os
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import jwt
from typing import Optional

from app.core.config import settings
from app.api.routes.videos import router as videos_router
from app.api.routes.auth import router as auth_router
from app.api.routes.processing import router as processing_router
from app.api.routes.users import router as users_router
from app.api.routes.health import router as health_router
from app.db.base import create_tables

# Create database tables
create_tables()

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

# Add security scheme
security = HTTPBearer(auto_error=False)

# Add authentication dependency
async def verify_token(request: Request, credentials: Optional[HTTPAuthorizationCredentials] = Depends(security)):
    # First check for HF Space JWT token in query params
    hf_token = request.query_params.get("__sign")
    if hf_token:
        try:
            # Verify the JWT token (you may want to add more verification)
            jwt.decode(hf_token, options={"verify_signature": False})
            return hf_token
        except jwt.InvalidTokenError:
            raise HTTPException(status_code=403, detail="Invalid HF Space token")
    
    # Then check for Bearer token
    if credentials:
        token = credentials.credentials
        if not token:
            raise HTTPException(
                status_code=401,
                detail="Invalid authentication credentials",
                headers={"WWW-Authenticate": "Bearer"},
            )
        return token
    
    # If no token found at all
    raise HTTPException(
        status_code=401,
        detail="Authentication required",
        headers={"WWW-Authenticate": "Bearer"},
    )

# Create FastAPI app
app = FastAPI(
    title="Behavior Analytics API",
    description="API for behavior analytics processing",
    version="1.0.0",
    docs_url="/docs",  # Swagger UI
    redoc_url="/redoc",  # ReDoc UI
    # Disable automatic redirect for trailing slashes
    redirect_slashes=False
)

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins for testing
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static files
app.mount("/uploads", StaticFiles(directory=str(settings.UPLOAD_DIR)), name="uploads")

# Create main API router
api_router = APIRouter()

# Include all routers except health
api_router.include_router(videos_router)
api_router.include_router(processing_router)
api_router.include_router(users_router)
api_router.include_router(auth_router)

# Include API router in app
app.include_router(api_router, prefix=settings.API_V1_STR)

# Mount health router directly to app
app.include_router(health_router)

@app.middleware("http")
async def log_requests(request: Request, call_next):
    """Log all requests and their processing time"""
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    logger.info(f"Path: {request.url.path} Method: {request.method} Time: {process_time:.2f}s Status: {response.status_code}")
    return response

@app.get("/")
async def root(token: str = Depends(verify_token)):
    """Root endpoint that returns API status"""
    return {
        "status": "ok",
        "message": "Behavior Analytics API is running",
        "version": "1.0.0"
    }

@app.get("/health")
async def health_check():
    """Health check endpoint for Hugging Face Spaces"""
    return {"status": "ok"}

@app.on_event("shutdown")
async def shutdown_event():
    loop = asyncio.get_running_loop()
    with ThreadPoolExecutor() as pool:
        await loop.run_in_executor(pool, shutdown_tasks)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860, reload=True)