from contextlib import asynccontextmanager

from fastapi import Depends, FastAPI, Header, HTTPException, status

from app.config import Settings, get_settings
from app.engine import InferenceEngine
from app.schemas import GenerateRequest, GenerateResponse, HealthResponse

engine: InferenceEngine | None = None


def verify_api_key(
    settings: Settings = Depends(get_settings), x_api_key: str | None = Header(default=None)
) -> None:
    if settings.api_key and x_api_key != settings.api_key:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid API key",
        )


@asynccontextmanager
async def lifespan(_: FastAPI):
    global engine
    settings = get_settings()
    engine = InferenceEngine(settings)
    yield
    if engine is not None:
        engine.close()
    engine = None


app = FastAPI(title="ROCm vLLM Inference API", version="1.0.0", lifespan=lifespan)


@app.get("/health", response_model=HealthResponse)
def health(settings: Settings = Depends(get_settings)) -> HealthResponse:
    return HealthResponse(status="ok", model=settings.served_model_name or settings.model_name)


@app.post("/v1/generate", response_model=GenerateResponse, dependencies=[Depends(verify_api_key)])
def generate(req: GenerateRequest, settings: Settings = Depends(get_settings)) -> GenerateResponse:
    if engine is None:
        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Engine not ready")
    if req.max_tokens > settings.max_tokens:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=f"max_tokens must be <= {settings.max_tokens}",
        )
    return engine.generate(req)