from __future__ import annotations

import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

from bucklake_ai.config import get_settings
from bucklake_ai.errors import BuckLakeAIError, InvalidScalerArtifact, ModelArtifactNotFound, ScalerArtifactNotFound
from bucklake_ai.schemas import (
    HealthResponse,
    ModelMetaResponse,
    PredictRequest,
    PredictResponse,
)
from bucklake_ai.service import InferenceService

settings = get_settings()
service = InferenceService(settings)


@asynccontextmanager
async def lifespan(_: FastAPI):
    logging.basicConfig(level=logging.INFO)
    service.preload()
    yield


app = FastAPI(title=settings.app_name, version=settings.app_version, lifespan=lifespan)


@app.exception_handler(ModelArtifactNotFound)
async def handle_model_not_found(_: Request, exc: ModelArtifactNotFound):
    return JSONResponse(
        status_code=503,
        content={"error": "model_artifact_not_found", "detail": str(exc)},
    )


@app.exception_handler(ScalerArtifactNotFound)
async def handle_scaler_not_found(_: Request, exc: ScalerArtifactNotFound):
    return JSONResponse(
        status_code=503,
        content={"error": "scaler_artifact_not_found", "detail": str(exc)},
    )


@app.exception_handler(InvalidScalerArtifact)
async def handle_invalid_scaler(_: Request, exc: InvalidScalerArtifact):
    return JSONResponse(
        status_code=503,
        content={"error": "invalid_scaler_artifact", "detail": str(exc)},
    )


@app.exception_handler(BuckLakeAIError)
async def handle_bucklake_error(_: Request, exc: BuckLakeAIError):
    return JSONResponse(
        status_code=400, content={"error": "bucklake_ai_error", "detail": str(exc)}
    )


@app.exception_handler(Exception)
async def handle_unexpected_error(_: Request, exc: Exception):
    return JSONResponse(
        status_code=500, content={"error": "internal_error", "detail": str(exc)}
    )


@app.get("/", response_model=ModelMetaResponse)
async def root():
    return service.get_model_meta()


@app.get("/health", response_model=HealthResponse)
async def health():
    return service.get_health()


@app.get("/api/model/meta", response_model=ModelMetaResponse)
async def model_meta():
    return service.get_model_meta()


@app.post("/api/predict/v2", response_model=PredictResponse)
async def predict_v2(request: PredictRequest):
    return service.predict(request)