"""
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,                     # fallback free-form policy
                   "preset":  "balanced"|"precision",
                   "categories": [                     # OPTIONAL structured rules
                       {"name": str,
                        "violation_rule": str,          # "flag this"
                        "not_violation_rule": str,      # "don't flag this"  (optional)
                        "exception_override": str},     # "exception"        (optional)
                       ...
                   ]
               }}
Response JSON: {"flagged": bool, "max_score": float,
                "categories": [{"category": str, "score": float, "flagged": bool}, ...]}

If `categories` is present and non-empty, each becomes its own PolicyLM Category
(so exceptions are read as exceptions and every rule gets its own score). Otherwise
the whole `policy` string is wrapped as a single catch-all category (legacy).
"""

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)
        raw_categories = inputs.get("categories") or []

        if raw_categories:
            # Structured rules: one PolicyLM Category per rule. Pass not_violation_rule
            # / exception_override only when supplied so we use the library defaults
            # otherwise.
            cats = []
            for i, c in enumerate(raw_categories, start=1):
                name = (c.get("name") or f"Rule {i}").strip() or f"Rule {i}"
                kwargs = {"violation_rule": c.get("violation_rule", "") or ""}
                not_violation = (c.get("not_violation_rule") or "").strip()
                exception = (c.get("exception_override") or "").strip()
                if not_violation:
                    kwargs["not_violation_rule"] = not_violation
                if exception:
                    kwargs["exception_override"] = exception
                cats.append(Category(name, **kwargs))
            policy = Policy(cats)
        else:
            # Legacy fallback: wrap the whole free-form document as one catch-all
            # category (loses per-rule scores and reads exceptions as violations).
            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,
        }