"""
Custom HF Inference Endpoint handler for musubilabs/policylm-1.7b.

PolicyLM is NOT OpenAI/vLLM-compatible — it needs its own `policylm_infer` helper
library and custom special tokens. So instead of a plain vLLM endpoint, deploy it
as a dedicated HF Inference Endpoint with this custom handler.

DEPLOYMENT
----------
1. Duplicate (or clone) the model repo so you can add files to it:
       huggingface-cli repo create policylm-1.7b-endpoint --type model   # under roosttools
       git clone https://huggingface.co/musubilabs/policylm-1.7b
   ...then push the model weights + policylm_infer.py into your own repo, OR add
   the two files below directly to a repo you control.
2. Add THIS file to that repo as `handler.py` (HF looks for `handler.py`).
3. Add a `requirements.txt` to that repo with whatever policylm_infer needs
   (at minimum: torch, transformers). policylm_infer.py must be in the repo.
4. Create an Inference Endpoint from that repo (GPU small is plenty for 1.7B),
   "Custom" container task.
5. Copy the endpoint URL and a token into the Space secrets:
       POLICYLM_ENDPOINT_URL      = https://xxxx.endpoints.huggingface.cloud
       POLICYLM_INFERENCE_TOKEN   = hf_...   (token with access to the endpoint)

REQUEST / RESPONSE CONTRACT (matches utils/model_interface.run_policylm)
------------------------------------------------------------------------
Request JSON:  {"inputs": {"message": str, "policy": str, "preset": "balanced"|"precision"}}
Response JSON: {"flagged": bool, "max_score": float,
                "categories": [{"category": str, "score": float, "flagged": bool}, ...]}
"""

from policylm_infer import Category, Policy, PolicyLM

# Flagging cutoffs for custom policies, from the model card.
CUTOFFS = {"precision": 0.335, "balanced": 0.275}
DEFAULT_CUTOFF = CUTOFFS["balanced"]


class EndpointHandler:
    def __init__(self, path: str = ""):
        # `path` is the local snapshot dir of the endpoint's model repo.
        self.model = PolicyLM.from_pretrained(path or "musubilabs/policylm-1.7b")

    def __call__(self, data: dict) -> dict:
        inputs = data.get("inputs", data) or {}
        message = inputs.get("message", "") or ""
        policy_text = inputs.get("policy", "") or ""
        preset = inputs.get("preset", "balanced")
        cutoff = CUTOFFS.get(preset, DEFAULT_CUTOFF)

        # The app supplies a single free-form policy document rather than discrete
        # categories, so wrap the whole document as one catch-all category. If you
        # later parse the policy into real categories, build them here instead.
        policy = Policy([Category("Policy violation", violation_rule=policy_text)])

        r = self.model.classify(message, policy, cutoff=cutoff)

        categories = [
            {"category": name, "score": float(c.score), "flagged": float(c.score) >= cutoff}
            for name, c in r.categories.items()
        ]
        max_score = max((c["score"] for c in categories), default=0.0)

        return {
            "flagged": bool(r.flagged),
            "max_score": max_score,
            "categories": categories,
        }