from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from predictor import load_benchmark, load_teacher, load_student

_models: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    _models["benchmark"] = load_benchmark()
    _models["teacher"]   = load_teacher()
    _models["student"]   = load_student()
    yield
    _models.clear()


app = FastAPI(title="AMANE Toxicity API", version="1.0.0", lifespan=lifespan)


class PredictRequest(BaseModel):
    text: str
    model: str = "student"


class BatchRequest(BaseModel):
    texts: list[str]
    model: str = "student"


@app.get("/health")
def health():
    return {
        name: {"loaded": p.loaded, "error": p.error}
        for name, p in _models.items()
    }


@app.post("/predict")
def predict(req: PredictRequest):
    p = _models.get(req.model)
    if p is None:
        raise HTTPException(status_code=404, detail=f"Unknown model: {req.model}")
    if not p.loaded:
        raise HTTPException(status_code=503, detail=f"Model not available: {p.error}")
    return p.predict(req.text)


@app.post("/predict_batch")
def predict_batch(req: BatchRequest):
    p = _models.get(req.model)
    if p is None:
        raise HTTPException(status_code=404, detail=f"Unknown model: {req.model}")
    if not p.loaded:
        raise HTTPException(status_code=503, detail=f"Model not available: {p.error}")
    return p.predict_batch(req.texts)