app/main.py

"""
Main FastAPI application for AUTH Microservice.
"""
import time
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.exceptions import RequestValidationError
from pydantic import ValidationError, BaseModel
from typing import Optional, List, Dict, Any
from jose import JWTError
from pymongo.errors import PyMongoError, ConnectionFailure, OperationFailure
from app.core.config import settings
from app.core.logging import setup_logging, get_logger

from app.nosql import connect_to_mongo, close_mongo_connection
from app.auth.controllers.router import router as auth_router
from app.auth.controllers.staff_router import router as staff_router
from app.auth.controllers.customer_router import router as customer_router
from app.auth.controllers.service_partners_router import router as service_partners_router
from app.internal.router import router as internal_router
from app.middleware.correlation_id import CorrelationIdMiddleware, get_correlation_id

# Setup logging
setup_logging(log_level=settings.LOG_LEVEL if hasattr(settings, "LOG_LEVEL") else "INFO")

logger = get_logger(__name__)


# Standard error response models
class ErrorDetail(BaseModel):
    """Detailed error information"""
    field: Optional[str] = None
    message: str
    type: Optional[str] = None


class ErrorResponse(BaseModel):
    """Standard error response format"""
    success: bool = False
    error: str
    detail: str
    errors: Optional[List[ErrorDetail]] = None
    request_id: Optional[str] = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage application lifespan events"""
    # Startup
    logger.info("Service starting", extra={"event": "service_starting"})
    await connect_to_mongo()
    
    # Initialize customer authentication collections
    try:
        from app.auth.db_init_customer import init_customer_auth_collections
        await init_customer_auth_collections()
        logger.info("Customer auth collections initialized", extra={"event": "db_init_success"})
    except Exception:
        logger.error(
            "Customer auth collection initialization failed",
            exc_info=True,
            extra={"event": "db_init_failure"},
        )
    
    logger.info("Service ready", extra={"event": "service_ready"})
    
    yield
    
    # Shutdown
    logger.info("Service stopping", extra={"event": "service_stopping"})
    await close_mongo_connection()
    logger.info("Service stopped", extra={"event": "service_stopped"})


import os

# Create FastAPI app
app = FastAPI(
    title="AUTH Microservice",
    description="Authentication & Authorization System - User Management, Login, JWT Tokens & Security",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
    lifespan=lifespan,
    root_path=os.getenv("ROOT_PATH", "")
)

# Add correlation ID middleware (must be added before other middlewares)
app.add_middleware(CorrelationIdMiddleware)


# Request logging middleware
@app.middleware("http")
async def log_requests(request: Request, call_next):
    """Log all incoming requests and responses with timing."""
    correlation_id = get_correlation_id()
    start_time = time.time()
    
    # Process request
    try:
        response = await call_next(request)
        
        # Calculate processing time
        duration_ms = round((time.time() - start_time) * 1000, 2)
        
        # Log response
        logger.info(
            "HTTP request",
            extra={
                "event": "http_request",
                "correlation_id": correlation_id,
                "method": request.method,
                "path": request.url.path,
                "status_code": response.status_code,
                "duration_ms": duration_ms,
                "client": request.client.host if request.client else None,
                "user_agent": request.headers.get("User-Agent", "")[:100],
            }
        )
        
        # Add custom headers
        response.headers["X-Process-Time"] = str(round(duration_ms / 1000, 3))
        
        return response
        
    except Exception:
        duration_ms = round((time.time() - start_time) * 1000, 2)
        logger.error(
            "HTTP request failed",
            exc_info=True,
            extra={
                "event": "http_request_failure",
                "correlation_id": correlation_id,
                "method": request.method,
                "path": request.url.path,
                "duration_ms": duration_ms,
            }
        )
        raise


# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.CORS_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Global exception handlers
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    """Handle request validation errors with detailed field information."""
    errors = []
    for error in exc.errors():
        field = " -> ".join(str(loc) for loc in error["loc"])
        errors.append({
            "field": field,
            "message": error["msg"],
            "type": error["type"]
        })
    
    logger.warning(
        "Request validation failed",
        extra={"event": "request_validation_failure", "path": request.url.path, "errors": errors},
    )
    
    return JSONResponse(
        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
        content={
            "success": False,
            "error": "Validation Error",
            "detail": "The request contains invalid data",
            "errors": errors
        }
    )


@app.exception_handler(ValidationError)
async def pydantic_validation_exception_handler(request: Request, exc: ValidationError):
    """Handle Pydantic validation errors."""
    logger.warning(
        "Pydantic validation failed",
        extra={"event": "request_validation_failure", "path": request.url.path, "errors": exc.errors()},
    )
    
    return JSONResponse(
        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
        content={
            "success": False,
            "error": "Data Validation Error",
            "detail": str(exc),
            "errors": exc.errors()
        }
    )


@app.exception_handler(JWTError)
async def jwt_exception_handler(request: Request, exc: JWTError):
    """Handle JWT token errors."""
    logger.warning(
        "JWT error",
        extra={"event": "token_invalid", "path": request.url.path, "error": str(exc)},
    )
    
    return JSONResponse(
        status_code=status.HTTP_401_UNAUTHORIZED,
        content={
            "success": False,
            "error": "Authentication Error",
            "detail": "Invalid or expired token",
            "headers": {"WWW-Authenticate": "Bearer"}
        }
    )


@app.exception_handler(PyMongoError)
async def mongodb_exception_handler(request: Request, exc: PyMongoError):
    """Handle MongoDB errors."""
    logger.error(
        "MongoDB request error",
        exc_info=True,
        extra={"event": "mongodb_operation_failure", "path": request.url.path, "error": str(exc)},
    )
    
    if isinstance(exc, ConnectionFailure):
        return JSONResponse(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            content={
                "success": False,
                "error": "Database Connection Error",
                "detail": "Unable to connect to the database. Please try again later."
            }
        )
    elif isinstance(exc, OperationFailure):
        return JSONResponse(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            content={
                "success": False,
                "error": "Database Operation Error",
                "detail": "A database operation failed. Please contact support if the issue persists."
            }
        )
    else:
        return JSONResponse(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            content={
                "success": False,
                "error": "Database Error",
                "detail": "An unexpected database error occurred."
            }
        )


@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    """Handle all uncaught exceptions."""
    correlation_id = get_correlation_id()
    
    logger.error(
        "Unhandled exception",
        exc_info=True,
        extra={
            "event": "unhandled_exception",
            "correlation_id": correlation_id,
            "method": request.method,
            "path": request.url.path,
            "client": request.client.host if request.client else None
        }
    )
    
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={
            "success": False,
            "error": "Internal Server Error",
            "detail": "An unexpected error occurred. Please try again later.",
            "correlation_id": correlation_id
        }
    )


# Health check endpoint
@app.get("/health", tags=["health"])
async def health_check():
    """
    Health check endpoint.
    Returns the service status and version.
    """
    try:
        return {
            "status": "healthy",
            "service": "auth-microservice",
            "version": "1.0.0"
        }
    except Exception as e:
        logger.error("Health check failed", exc_info=True, extra={"event": "health_check_failure"})
        return JSONResponse(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            content={
                "status": "unhealthy",
                "service": "auth-microservice",
                "error": str(e)
            }
        )


# Debug endpoint to check database status
@app.get("/debug/db-status", tags=["debug"])
async def check_db_status():
    """
    Check database connection and user count.
    Returns information about database collections and sample data.
    
    Raises:
        HTTPException: 500 - Database connection error
    """
    try:
        from app.nosql import get_database
        from app.constants.collections import AUTH_SYSTEM_USERS_COLLECTION, AUTH_ACCESS_ROLES_COLLECTION, SCM_ACCESS_ROLES_COLLECTION
        
        db = get_database()
        
        if db is None:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="Database connection not available"
            )
        
        users_count = await db[AUTH_SYSTEM_USERS_COLLECTION].count_documents({})
        roles_count = await db[AUTH_ACCESS_ROLES_COLLECTION].count_documents({})
        scm_roles_count = await db[SCM_ACCESS_ROLES_COLLECTION].count_documents({})
        
        # Get sample user to verify
        sample_user = await db[AUTH_SYSTEM_USERS_COLLECTION].find_one(
            {"email": "superadmin@cuatrolabs.com"},
            {"email": 1, "username": 1, "role": 1, "status": 1, "_id": 0}
        )
        
        return {
            "status": "connected",
            "database": settings.MONGODB_DB_NAME,
            "collections": {
                "users": users_count,
                "auth_roles": roles_count,
                "scm_roles": scm_roles_count
            },
            "superadmin_exists": sample_user is not None,
            "sample_user": sample_user if sample_user else None
        }
    except HTTPException:
        raise
    except Exception as e:
        logger.error("Database status check failed", exc_info=True, extra={"event": "db_status_failure"})
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Failed to check database status: {str(e)}"
        )


# Include routers with new organization
app.include_router(auth_router)                    # /login, /refresh, /me, /logout, /users/* - All auth & user management
app.include_router(staff_router)                   # /staff/* - Staff authentication (mobile OTP)
app.include_router(customer_router)                # /customer/* - Customer authentication (OTP)
app.include_router(service_partners_router)    # /service-professional/* - Service professional authentication (OTP)
app.include_router(internal_router)                # /internal/* - Internal API endpoints


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=8002,
        reload=True,
        log_level="info"
    )