"""
zybrAI ML inference service — deploy to a FREE Hugging Face Space (Docker SDK).

Hosts the two classifiers (DeBERTa prompt-injection + toxic-RoBERTa) and serves
them in the SAME request/response shape as the old HF serverless Inference API,
so zybrAI's ml_engine can point at this Space with zero code change:

  POST /models/{model_id}      body: {"inputs": "<text>"}
  -> [{"label": "...", "score": 0.98}, ...]   (text-classification, all scores)

zybrAI backend then sets:
  ML_MODE   = api
  HF_API_BASE = https://<your-space-subdomain>.hf.space
  HF_API_KEY  = <any value; only used if SPACE_SECRET is set below>

Free HF Space (CPU basic: 2 vCPU / 16GB RAM) comfortably runs both models.
First request per model cold-loads it (~10-30s); subsequent calls are fast.
"""
import os
from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel
from transformers import pipeline

# Models this Space hosts. Add more here if zybrAI adds ML detectors.
HOSTED = {
    "protectai/deberta-v3-base-prompt-injection-v2",
    "unitary/unbiased-toxic-roberta",
}
_PIPES: dict = {}

# Optional shared secret so randoms can't abuse your free Space. Set SPACE_SECRET
# as a Space "Secret" and the same value as HF_API_KEY on the zybrAI backend.
SPACE_SECRET = os.getenv("SPACE_SECRET", "")

app = FastAPI(title="zybrAI ML inference", version="1.0")


class Inp(BaseModel):
    inputs: str


def _pipe(model_id: str):
    if model_id not in _PIPES:
        _PIPES[model_id] = pipeline(
            "text-classification", model=model_id, top_k=None, truncation=True
        )
    return _PIPES[model_id]


@app.get("/")
def health():
    return {"status": "ok", "hosted": sorted(HOSTED), "loaded": sorted(_PIPES.keys())}


@app.post("/models/{model_id:path}")
def infer(model_id: str, body: Inp, authorization: str = Header(default="")):
    if SPACE_SECRET:
        token = authorization.replace("Bearer ", "").strip()
        if token != SPACE_SECRET:
            raise HTTPException(status_code=401, detail="unauthorized")
    if model_id not in HOSTED:
        raise HTTPException(status_code=404, detail=f"model {model_id} not hosted")
    result = _pipe(model_id)(body.inputs[:512])
    # top_k=None returns list[dict] for a single input (or list[list[dict]]).
    items = result[0] if result and isinstance(result[0], list) else result
    return items  # zybrAI's parser accepts a flat [{label,score},...] list