app/main.py

# app/main.py
"""
Application entrypoint.

- Configures FastAPI app and CORS middleware.
- Includes embed router.
- Starts uvicorn when run as __main__.
"""

import torch
from contextlib import asynccontextmanager
import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from .config import settings
from .logger import logger
from .controllers.embed_controller import router as embed_router, service as embedding_service

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup logic
    gpu_available = torch.cuda.is_available()
    device_to_use = settings.DEVICE or ("cuda" if gpu_available else "cpu")
    
    logger.info("┌────────────────────────────────────────┐")
    logger.info("│          Embedding Server Start        │")
    logger.info("└────────────────────────────────────────┘")
    
    if gpu_available:
        logger.info(f"GPU Detected: [bold]{torch.cuda.get_device_name(0)}[/bold]")
    else:
        logger.warning("GPU not detected, falling back to CPU.")
        
    logger.info(f"Active Device: {device_to_use}")
    logger.info(f"Serving on: {settings.HOST}:{settings.PORT}")

    # Eagerly load the model to RAM/VRAM on startup
    try:
        await embedding_service.load_model()
    except Exception as e:
        logger.error("CRITICAL: Failed to load model during startup.")
        logger.exception(e)
        # We don't exit here to allow the health check to remain accessible, 
        # but the model status will be 'unloaded'.
    
    yield
    
    # Shutdown logic
    logger.info("Shutting down server...")

app = FastAPI(
    title="BGE-M3 Multilingual Embedding Server",
    summary="High-performance API for BAAI/bge-m3 embeddings.",
    description="""
        ## Features
        - **Multilingual Support**: Embed text in 100+ languages.
        - **GPU Acceleration**: Automatically uses CUDA if available.
        - **Batch Processing**: Efficiently handle lists of text.
        - **Security**: Strict API Key protection.

        ## Usage
        - **Required Header**: `x-api-key` must be provided for all `POST` / `GET` requests.
        - **Interactive Docs**: Use the **Try it out** button to test endpoints directly from the browser.
        """,
    version="1.0.0",
    lifespan=lifespan,
    docs_url="/docs",  
    redoc_url="/redoc" 
)

# Global Exception Handler for unhandled errors
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error(f"Unhandled Exception: {type(exc).__name__}")
    logger.error(f"Path: {request.url.path}")
    logger.exception(exc) # This will log the full traceback in red
    return JSONResponse(
        status_code=500,
        content={"detail": "An internal server error occurred.", "type": type(exc).__name__},
    )

# CORS middleware - allow all (security handled by API Key guard)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)

import time
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.perf_counter()
    response = await call_next(request)
    process_time = (time.perf_counter() - start_time) * 1000
    logger.info(f"Request: {request.method} {request.url.path} - Completed in {process_time:.2f}ms")
    return response

@app.get("/", tags=["status"])
async def root():
    """Root endpoint providing server status and documentation links."""
    return {
        "status": "online",
        "message": "BGE-M3 Multilingual Embedding Server is running.",
        "version": app.version,
        "docs_url": "/docs",
        "health_check": "/embed/health"
    }

app.include_router(embed_router)



if __name__ == "__main__":
    uvicorn.run("app.main:app", host=settings.HOST, port=settings.PORT, workers=settings.WORKERS)