"""API REST de prediction de consommation electrique (FastAPI).

Le service orchestre la chaine de prediction :
  1. recevoir la requete (client_id, date) ;
  2. recuperer les features du client (feature store simule) ;
  3. charger le modele promu (Model Registry, au demarrage) ;
  4. calculer la prediction et repondre en JSON.

Lancement : uvicorn lab.serving.api:app --host 0.0.0.0 --port 8000
Swagger UI : /docs
"""

import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException

from . import features, registry
from .schemas import (
    BatchPredictionRequest,
    BatchPredictionResponse,
    PredictionRequest,
    PredictionResponse,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Modele charge une fois au demarrage (couteux) et reutilise a chaque requete.
_state: dict[str, registry.LoadedModel] = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    _state["model"] = registry.load_champion()
    yield
    _state.clear()


app = FastAPI(
    title="Electricity Consumption Prediction API",
    description="Expose le modele promu (MLflow Model Registry) via une API REST.",
    version="1.0.0",
    lifespan=lifespan,
)


def _get_model() -> registry.LoadedModel:
    model = _state.get("model")
    if model is None:  # modele indisponible au demarrage
        raise HTTPException(status_code=503, detail="Modele non charge.")
    return model


def _predict(client_id: str, model: registry.LoadedModel) -> PredictionResponse:
    feats = features.get_features(client_id)  # peut lever UnknownClientError
    prediction = model.predict_one(feats)
    return PredictionResponse(
        client_id=client_id,
        prediction_kwh=prediction,
        model_name=model.name,
        model_version=model.version,
    )


@app.get("/health", summary="Verification de l'etat du service")
def health() -> dict[str, str]:
    """Endpoint de sante : renvoie 200 si le service repond."""
    return {"status": "ok"}


@app.post("/predict", response_model=PredictionResponse, summary="Prediction unitaire")
def predict(request: PredictionRequest) -> PredictionResponse:
    model = _get_model()
    try:
        return _predict(request.client_id, model)
    except features.UnknownClientError:
        # Partie 3 : client inconnu -> erreur cliente, pas un plantage du service.
        raise HTTPException(
            status_code=404,
            detail=f"Client inconnu : aucune feature pour '{request.client_id}'.",
        )


@app.post(
    "/predict/batch",
    response_model=BatchPredictionResponse,
    summary="Prediction en batch (plusieurs clients)",
)
def predict_batch(request: BatchPredictionRequest) -> BatchPredictionResponse:
    model = _get_model()
    predictions: list[PredictionResponse] = []
    unknown: list[str] = []
    for client_id in request.client_ids:
        try:
            predictions.append(_predict(client_id, model))
        except features.UnknownClientError:
            unknown.append(client_id)
    return BatchPredictionResponse(predictions=predictions, unknown_client_ids=unknown)