#!/usr/bin/env python3
import sys, json, argparse
import numpy as np
import joblib
import onnxruntime as ort
from transformers import AutoTokenizer
from pathlib import Path

DIR = Path("/usr/local/nagiosxi/var/nagios/xihelper")

ENC_DIR = DIR / "all-MiniLM-L6-v2-onnx"
ONNX_PATH = ENC_DIR / "model.onnx"

tokenizer = AutoTokenizer.from_pretrained(str(ENC_DIR), local_files_only=True)
session = ort.InferenceSession(str(ONNX_PATH), providers=["CPUExecutionProvider"])

CLASSIFIERS = {
    "action": {
        "clf": DIR / "logreg_head_action.joblib",
        "labels": "/usr/local/nagiosxi/html/includes/classifiers/labels_action.json",
    },
    "hosttype": {
        "clf": DIR / "logreg_head_hosttype.joblib",
        "labels": "/usr/local/nagiosxi/html/includes/classifiers/labels_hosttype.json",
    },
}

#  Encoder 
def encode(text: str, max_length: int = 128) -> np.ndarray:
    enc = tokenizer([text], padding="max_length", truncation=True,
                    max_length=max_length, return_tensors="np")
    allowed = {i.name for i in session.get_inputs()}
    ort_inputs = {k: v for k, v in enc.items() if k in allowed}
    emb = session.run(None, ort_inputs)[0]
    return emb.astype(np.float32)

parser = argparse.ArgumentParser()
parser.add_argument("text")
parser.add_argument("--action", action="store_true")
parser.add_argument("--hosttype", action="store_true")
args = parser.parse_args()

if not (args.action or args.hosttype):
    print(json.dumps({"error": "must specify --action or --hosttype"}))
    sys.exit(1)

task = "action" if args.action else "hosttype"

clf = joblib.load(CLASSIFIERS[task]["clf"])
labels = json.load(open(CLASSIFIERS[task]["labels"]))

emb = encode(args.text.strip())
probs = clf.predict_proba(emb)[0]
idx = int(np.argmax(probs))

print(json.dumps({"label": labels[idx],"score": float(probs[idx])}))
