app/server.py

import os
import asyncio
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware

# Import our modules
from app.config import config
from app.ml_utils import model_manager
from app.database import db_manager
from app.cloudinary_utils import initialize_cloudinary

# Import route modules
from app.routes.prediction import router as prediction_router
from app.routes.history import router as history_router
from app.routes.system import router as system_router


class ConnectionErrorMiddleware(BaseHTTPMiddleware):
    """Handle connection errors gracefully for Hugging Face deployment"""
    async def dispatch(self, request: Request, call_next):
        try:
            response = await call_next(request)
            return response
        except ConnectionResetError as e:
            print(f"Connection reset error for {request.url}: {e}")
            return JSONResponse(
                status_code=499,
                content={"detail": "Connection was closed by client"}
            )
        except asyncio.CancelledError:
            print(f"Request cancelled for {request.url}")
            return JSONResponse(
                status_code=499,
                content={"detail": "Request was cancelled"}
            )
        except Exception as e:
            print(f"Unexpected error for {request.url}: {e}")
            return JSONResponse(
                status_code=500,
                content={"detail": "Internal server error"}
            )


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager for startup and shutdown"""
    print("===== Application Startup =====")
    
    try:
        # Configure asyncio for Windows if needed
        if os.name == 'nt':
            asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
        
        # Validate essential configuration
        config.validate_required_config()
        print("✓ Configuration validation passed")
        
        # Load ML model (essential)
        if model_manager.load_model():
            print("✓ Model loaded successfully")
        else:
            print("✗ Model loading failed")
            raise RuntimeError("Model loading failed - cannot start application")
        
        # Initialize optional services
        await db_manager.initialize()
        if db_manager.is_available():
            print("✓ Database initialized")
        else:
            print("ℹ Database not available - predictions will not be persisted")
        
        initialize_cloudinary()
        print("ℹ Cloudinary initialization completed")
        
        print(f"✓ Server ready on {config.HOST}:{config.PORT}")
        print("===== Startup Complete =====")
        
    except Exception as e:
        print(f"✗ Startup error: {e}")
        raise
    
    yield  # Application runs here
    
    print("===== Application Shutdown =====")


# Initialize FastAPI app with lifespan manager
app = FastAPI(
    title="Brain Stroke Segmentation API",
    description="AI-powered brain stroke detection and segmentation",
    version="1.0.0",
    debug=config.DEBUG,
    lifespan=lifespan
)

# Add middlewares
app.add_middleware(ConnectionErrorMiddleware)

# CORS middleware - more permissive for Hugging Face
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"] if "*" in config.CORS_ORIGINS else config.CORS_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include routers
app.include_router(prediction_router, tags=["Prediction"])
app.include_router(history_router, tags=["History"])
app.include_router(system_router, tags=["System"])


if __name__ == "__main__":
    import uvicorn
    
    try:
        print(f"Starting server on {config.HOST}:{config.PORT}")
        uvicorn.run(
            "app.server:app",
            host=config.HOST,
            port=config.PORT,
            reload=config.DEBUG,
            access_log=False,
            server_header=False,
            date_header=False,
        )
    except KeyboardInterrupt:
        print("Server stopped by user")
    except Exception as e:
        print(f"Server startup failed: {e}")
        raise