# -*- coding: utf-8 -*-
"""
serve.py — FastAPI endpoint สำหรับ Pathumma Cross-border Guardrail

    pip install fastapi uvicorn openthai-systemone
    GUARDRAIL_MODEL=TODO-org/pathumma-crossborder-guardrail uvicorn serve:app --host 0.0.0.0 --port 8080

Endpoints
  POST /v1/guardrail/check         {"text": "..."}                       -> Decision (native)
  POST /v1/guardrail/check_batch   {"texts": ["...", "..."]}             -> [Decision]
  POST /v1/moderations             OpenAI-compatible  {"input": "..." | ["..."]}
                                   -> {"id","model","results":[{"flagged","categories","category_scores","routing",...}]}
  GET  /healthz

/v1/moderations ใช้ต่อกับ framework ที่รับ OpenAI moderation schema (LiteLLM openai_moderation guardrail,
Open WebUI, NeMo Guardrails, API gateway) โดยชี้ api_base มาที่ server นี้
  flagged = routing != "foreign_ok"   (คือ "ห้ามส่งออกต่างประเทศแบบไม่มีเงื่อนไข")
  categories = {<7 หมวด>: bool, "needs_review": bool, "routing_conditional": bool, "routing_sovereign_only": bool}
"""
import os, time, uuid, logging
from typing import List, Union
from fastapi import FastAPI
from pydantic import BaseModel
from guardrail import Guardrail

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
MODEL = os.environ.get("GUARDRAIL_MODEL", "./systemone-guardrail-ft")
MODEL_NAME = os.environ.get("GUARDRAIL_MODEL_NAME", "pathumma-crossborder-guardrail")
app = FastAPI(title="Pathumma Cross-border Guardrail", version="1.1.0")
g = Guardrail(MODEL)
CATS = list(g.cfg["criteria"].keys())

class CheckReq(BaseModel):  text: str
class BatchReq(BaseModel):  texts: List[str]
class ModerationReq(BaseModel):
    input: Union[str, List[str]]
    model: str | None = None

def _log(d: dict):
    logging.info("decision cat=%s routing=%s conf=%.3f review=%s src=%s ms=%.0f",
                 d["category"], d["routing"], d["confidence"], d["needs_review"], d["source"], d["latency_ms"])  # ไม่ log text

@app.get("/healthz")
def healthz(): return {"ok": True, "model": MODEL, "config_version": g.cfg.get("version")}

@app.post("/v1/guardrail/check")
def check(r: CheckReq):
    d = g.check(r.text).to_dict(); _log(d); return d

@app.post("/v1/guardrail/check_batch")
def check_batch(r: BatchReq):
    out = [d.to_dict() for d in g.check_batch(r.texts)]
    for d in out: _log(d)
    return out

def _to_moderation_result(d: dict) -> dict:
    scores = {c: float(d["probabilities"].get(c, 0.0)) for c in CATS}
    cats = {c: (d["category"] == c) for c in CATS}
    cats.update({"needs_review": d["needs_review"],
                 "routing_conditional": d["routing"] == "conditional",
                 "routing_sovereign_only": d["routing"] == "sovereign_only"})
    scores.update({"needs_review": 1.0 - d["confidence"],
                   "routing_conditional": float(d["routing"] == "conditional"),
                   "routing_sovereign_only": float(d["routing"] == "sovereign_only")})
    return {"flagged": d["routing"] != "foreign_ok",
            "categories": cats, "category_scores": scores,
            # ฟิลด์เสริม (นอก schema OpenAI แต่ client ส่วนใหญ่ ignore ได้)
            "category": d["category"], "routing": d["routing"], "confidence": d["confidence"],
            "abstain": d["abstain"], "needs_review": d["needs_review"], "source": d["source"]}

@app.post("/v1/moderations")
def moderations(r: ModerationReq):
    inputs = [r.input] if isinstance(r.input, str) else list(r.input)
    results = []
    for t in inputs:
        d = g.check(t).to_dict(); _log(d); results.append(_to_moderation_result(d))
    return {"id": f"modr-{uuid.uuid4().hex[:24]}", "model": r.model or MODEL_NAME,
            "created": int(time.time()), "results": results}