"""HTTP service around one Bauxi model, built to stay inside a 100m / 80 MiB pod."""

from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated, Any

import anyio
from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
from pydantic import BaseModel, ConfigDict, Field

from bauxi import MAX_CHARS, Bauxi, Label, __version__, release_unused_memory


class DecideRequest(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    decision: str = Field(min_length=1, max_length=MAX_CHARS)
    evidence: str = Field(min_length=1, max_length=MAX_CHARS)


class DecideResponse(BaseModel):
    prediction: Label
    probabilities: dict[Label, float]
    calibrated: bool


@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
    # One worker reuses one native arena; a thread per caller would grow one each.
    anyio.to_thread.current_default_thread_limiter().total_tokens = 1
    yield


app = FastAPI(title="Bauxi", version=__version__, lifespan=lifespan)
router = APIRouter()


def get_model(request: Request) -> Bauxi:
    model: Bauxi = request.app.state.model
    return model


@router.post("/decide")
def decide(body: DecideRequest, model: Annotated[Bauxi, Depends(get_model)]) -> DecideResponse:
    try:
        result = model.predict(body.decision, body.evidence)
    except ValueError as error:
        raise HTTPException(status_code=422, detail=str(error)) from error
    finally:
        release_unused_memory()
    return DecideResponse(
        prediction=result.label,
        probabilities=result.probabilities,
        calibrated=result.calibrated,
    )


@router.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok", "device": "cpu"}


@router.get("/runtime")
def runtime(model: Annotated[Bauxi, Depends(get_model)]) -> dict[str, Any]:
    """What is loaded and what the cgroup measures, for the verifier and benchmarks."""
    root = Path("/sys/fs/cgroup")
    names = ("memory.current", "memory.peak", "memory.max", "memory.events", "cpu.max", "cpu.stat")
    return {
        "version": model.metadata["version"],
        "artifact": model.artifact,
        "artifact_sha256": model.metadata["artifacts"][model.artifact]["sha256"],
        "providers": model.session.get_providers(),
        "cgroup": {
            name: (root / name).read_text().strip() for name in names if (root / name).exists()
        },
    }


app.include_router(router)