"""Standalone inference for the DistilBERT email-triage model (see module docstring)."""
INPUT_FORMAT = ("Subject: ...\nFrom: ***@domain\nReply-To: ...\nBody: ...\nURLs: ... "
                "(§4 normalize_row + PII redaction; truncated at max_length)")
import json, os
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_DIR = os.environ.get('JEV_MODEL_DIR', os.path.dirname(os.path.abspath(__file__)))
with open(os.path.join(MODEL_DIR, 'label_map.json')) as _f:
    _LM = json.load(_f)
LABELS = _LM['labels']
ACTION = {'legitimate': 'allow', 'spam': 'filter_to_spam',
          'phishing': 'quarantine_or_review', 'malware': 'quarantine_or_review'}
_THRESHOLD = float(_LM.get('review_threshold', 0.55))
_MAX = int(_LM.get('max_length', 512))
_tok = AutoTokenizer.from_pretrained(MODEL_DIR)
_model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR).eval()

@torch.no_grad()
def predict(texts, threshold=_THRESHOLD, batch_size=16):
    outs, all_p = [], []
    for i in range(0, len(texts), batch_size):
        t = _tok(list(texts[i:i + batch_size]), truncation=True, padding=True,
                 max_length=_MAX, return_tensors='pt')
        all_p.append(torch.softmax(_model(**t).logits, -1).cpu().numpy())
    import numpy as _np
    P = _np.concatenate(all_p) if all_p else _np.zeros((0, len(LABELS)))
    for t, p in zip(texts, P):
        scores = {l: float(p[LABELS.index(l)]) for l in LABELS}
        top = LABELS[int(p.argmax())]
        if not str(t).strip() or float(p.max()) < threshold:
            outs.append({'label': 'review', 'scores': scores, 'action': 'review'})
        else:
            outs.append({'label': top, 'scores': scores, 'action': ACTION[top]})
    return outs

if __name__ == '__main__':
    demo = ['Subject: Meeting moved to 2 PM\nFrom: \nReply-To: \nBody: See the calendar invite.\nURLs: ', '']
    print(json.dumps(predict(demo), indent=2))