"""Çözücülü tespit: metnin çözülmüş varyantlarını da sınıflandırır, en yüksek skoru döndürür.

  from detect import Detector
  d = Detector()                    # v2, eşik 0,05 (v2.1; 0,3 ilk yayındı)
  d = Detector(mode="recall")       # v1 + v2 ortalaması, eşik 0,2: daha çok yakalar, daha çok yanlış alarm
  d("Önceki mesajımdaki formatı unut, tablo yap.")
  # {'label': 'BENIGN', 'score': 0.0, 'decoded': [], 'max_from': None}

`score` INJECTION olasılığıdır (orijinal ve çözülmüş varyantların en yükseği). `decoded` çözülen varyantlar
(base64 / hex / Sezar / ters yazım / harf aralama / görünmez Unicode); `max_from`, en yüksek skoru veren varyant
orijinal değilse odur.
"""
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

try:
    from guard import variants
except ImportError:  # eğitim deposunda paket olarak
    from pinj.guard import variants

REPO = "halilneed/turkish-prompt-injection-classifier"
# Eşikler ve ağırlık, test bölmelerine dokunmayan seçim protokolünde seçildi (model kartı: "How v2 was chosen")
MODES = {"default": {"revisions": [None], "weights": [1.0], "threshold": 0.05},
         "recall": {"revisions": ["v1", None], "weights": [0.5, 0.5], "threshold": 0.2}}


class _Model:
    def __init__(self, path, revision, device):
        self.tok = AutoTokenizer.from_pretrained(path, revision=revision)
        self.model = AutoModelForSequenceClassification.from_pretrained(path, revision=revision).eval().to(device)
        self.idx = self.model.config.label2id["INJECTION"]
        self.device = device

    @torch.inference_mode()
    def __call__(self, texts):
        enc = self.tok(texts, truncation=True, max_length=512, padding=True, return_tensors="pt").to(self.device)
        return self.model(**enc).logits.float().softmax(-1)[:, self.idx].tolist()


class Detector:
    def __init__(self, model=REPO, mode="default", threshold=None, device=None, paths=None):
        """paths: yerel deneme için mod revizyonları yerine klasör listesi (ör. [v1_klasörü, v2_klasörü])."""
        cfg = MODES[mode]
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        if paths:
            self.models = [_Model(p, None, self.device) for p in paths]
        else:
            self.models = [_Model(model, r, self.device) for r in cfg["revisions"]]
        self.weights = cfg["weights"]
        self.threshold = cfg["threshold"] if threshold is None else threshold

    def scores(self, texts):
        per = [m(texts) for m in self.models]
        return [sum(w * p[i] for w, p in zip(self.weights, per)) for i in range(len(texts))]

    def __call__(self, text):
        vs = variants(text)
        sc = self.scores(vs)
        i = max(range(len(vs)), key=sc.__getitem__)
        return {"label": "INJECTION" if sc[i] >= self.threshold else "BENIGN", "score": round(sc[i], 4),
                "decoded": vs[1:], "max_from": vs[i] if i else None}