"""Prompt injection dedektörlerini Türkçe test setlerinde ölçer.

Saldırı setleri (Hub'daki açık lisanslı Türkçe setlerin test bölmeleri; eğitimde kullanılmaz):
  conv   3nesdeniz/turkish-conversation-prompt-injection  test 120  (CC-BY-4.0)
  1k     3nesdeniz/turkish-prompt-injection-1k             test 113  (CC-BY-4.0)
  berat  beratcmn/turkish-prompt-injections               test 115  (Apache-2.0)
  fevzi  fevziegeyurtsevenler/turkish-prompt-injection    tümü 107, yalnızca saldırı (CC-BY-4.0)
Yanlış alarm setleri (hepsi zararsız, eğitimde kullanılmayan CC0 Türkçe metinler):
  fp_kvkk   halilneed/turkish-kvkk-classification-benchmark   1.125 iş metni
  fp_bseby  halilneed/turkish-bseby-classification-benchmark  750 bankacılık metni
  fp_pii    cagrigungor/turkish-pii-masking-benchmark          1.000 talimat + metin

Kullanım:
  python -m pinj.evaluate --model runs/pinj-berturk/final --name ours --out results/pinj
  python -m pinj.evaluate --rule --name keyword-rule --out results/pinj
"""
import argparse
import csv
import json
import re
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

EXT = Path("data/external")
REL = Path("../release")
# fevzi setindeki yer tutucular sabit, açıkça sahte değerlerle doldurulur
FILL = {"{SECRET}": "sk-test-4f2a", "{URL}": "ornek-log.net", "{EMAIL}": "x@ornek-log.net"}


def jl(p):
    return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]


def fill(t):
    for k, v in FILL.items():
        t = t.replace(k, v)
    return t


def load_sets():
    s = {}
    s["conv"] = [{"text": r["text"], "label": r["label"], "family": r["attack_family"], "cat": r["category"]}
                 for r in jl(EXT / "3nesdeniz_conv.test.jsonl")]
    s["1k"] = [{"text": r["text"], "label": r["label"], "family": r["attack_family"], "cat": r["technique"]}
               for r in jl(EXT / "3nesdeniz_1k.test.jsonl")]
    s["berat"] = [{"text": r["text"], "label": r["label"]} for r in jl(EXT / "beratcmn.test.jsonl")]
    s["fevzi"] = [{"text": fill(r["payload"]), "label": 1, "family": r["category"]}
                  for r in jl(EXT / "fevzi.train.jsonl")]
    # Görünmez Unicode ile gizlenmiş talimatlar (42, CC-BY-4.0); gizli metin TR ya da EN, görünür metin masum
    s["invisible"] = [{"text": fill(r["full_text"]), "label": 1, "family": r["technique"]}
                      for r in jl(EXT / "fevzi_invisible.train.jsonl")]
    s["fp_kvkk"] = [{"text": r["text"], "label": 0}
                    for r in jl(REL / "turkish-kvkk-classification-benchmark/all_annotations.jsonl")]
    s["fp_bseby"] = [{"text": r["text"], "label": 0}
                     for r in jl(REL / "turkish-bseby-classification-benchmark/all_annotations.jsonl")]
    with open("data/benchmark/benchmark_1000.csv", encoding="utf-8-sig") as f:
        s["fp_pii"] = [{"text": f"{r['instruction']}\n\n{r['input']}", "label": 0} for r in csv.DictReader(f)]
    return s


def load_selection_sets(n_fp=800, seed=7):
    """Model seçimi için: test bölmelerine dokunmaz. Kaynak X dışarıda bırakılarak eğitilen model X'in
    train+validation bölmelerinde ölçülür; yanlış alarm, eğitimde kullanılmayan iç bölmelerde."""
    import random
    rng = random.Random(seed)
    s = {}
    s["sel_conv"] = [{"text": r["text"], "label": r["label"], "family": r["attack_family"], "cat": r["category"]}
                     for sp in ("train", "validation") for r in jl(EXT / f"3nesdeniz_conv.{sp}.jsonl")]
    s["sel_1k"] = [{"text": r["text"], "label": r["label"], "family": r["attack_family"], "cat": r["technique"]}
                   for sp in ("train", "validation") for r in jl(EXT / f"3nesdeniz_1k.{sp}.jsonl")]
    s["sel_berat"] = [{"text": r["text"], "label": r["label"]} for r in jl(EXT / "beratcmn.train.jsonl")]
    for name, path in (("fpdev_kvkk", "data/kvkk3/test.jsonl"), ("fpdev_bseby", "data/bseby/test.jsonl")):
        texts = sorted({r["text"] for r in jl(path)})
        rng.shuffle(texts)
        s[name] = [{"text": t, "label": 0} for t in texts[:n_fp]]
    v02 = jl("data/v02/dev.jsonl")
    rng.shuffle(v02)
    s["fpdev_pii"] = [{"text": r["instruction"] + "\n\n" + r["input"], "label": 0} for r in v02[:n_fp]]
    return s


# Anahtar sözcük kuralı: "küçük bir model yerine kural yeter mi" sorusunun taban çizgisi
RULE = re.compile(
    r"(önceki|yukarıdaki|tüm|bütün)\s+(talimat|kural|komut)\w*\s+(unut|yok\s*say|görmezden|dikkate\s*alma)"
    r"|sistem\s*(prompt|istem|talimat|mesaj)\w*|system\s*prompt|ignore\s+(all\s+)?previous"
    r"|jailbreak|\bDAN\b|geliştirici\s+mod|developer\s+mode|filtresiz|kısıtlama\w*\s+(yok|kaldır)"
    r"|kuralların\w*\s+(neler|ne)|gizli\s+(talimat|anahtar)|(yönetici|admin)\s+olarak\s+(izin|yetki)",
    re.IGNORECASE)


def rule_scores(texts):
    return np.array([1.0 if RULE.search(t.replace("İ", "i")) else 0.0 for t in texts])


def injection_index(model):
    for i, lb in model.config.id2label.items():
        if "inj" in str(lb).lower() or "jail" in str(lb).lower() or str(lb).upper() == "LABEL_1":
            return int(i)
    raise ValueError(f"injection etiketi bulunamadı: {model.config.id2label}")


@torch.no_grad()
def model_scores(model, tok, texts, max_len, bs=32):
    dev = next(model.parameters()).device
    idx = injection_index(model)
    order = np.argsort([len(t) for t in texts])
    out = np.zeros(len(texts))
    for i in range(0, len(texts), bs):
        sel = order[i:i + bs]
        enc = tok([texts[j] for j in sel], truncation=True, max_length=max_len, padding=True,
                  return_tensors="pt").to(dev)
        prob = model(**enc).logits.float().softmax(-1)[:, idx].cpu().numpy()
        out[sel] = prob
    return out


def guarded_scores(model, tok, texts, max_len):
    """Çözücüyle: her metnin varyantları skorlanır, en yüksek skor alınır."""
    from pinj.guard import variants
    flat, owner = [], []
    for i, t in enumerate(texts):
        for v in variants(t):
            flat.append(v)
            owner.append(i)
    sc = model_scores(model, tok, flat, max_len)
    out = np.zeros(len(texts))
    for i, v in zip(owner, sc):
        out[i] = max(out[i], v)
    return out


def metrics(rows, score, th):
    y = np.array([r["label"] for r in rows])
    p = (score >= th).astype(int)
    tp, fp = int(((p == 1) & (y == 1)).sum()), int(((p == 1) & (y == 0)).sum())
    fn, tn = int(((p == 0) & (y == 1)).sum()), int(((p == 0) & (y == 0)).sum())
    m = {"n": len(y), "n_attack": int(y.sum())}
    if y.sum():
        m["recall"] = round(tp / (tp + fn), 4)
    if (y == 0).sum():
        m["fpr"] = round(fp / (fp + tn), 4)
    if y.sum() and (y == 0).sum():
        pr = tp / (tp + fp) if tp + fp else 0.0
        rc = tp / (tp + fn)
        m["precision"] = round(pr, 4)
        m["f1"] = round(2 * pr * rc / (pr + rc) if pr + rc else 0.0, 4)
        m["accuracy"] = round((tp + tn) / len(y), 4)
    # Dilimler: saldırıda ailenin yakalanma oranı, zararsızda kategorinin yanlış alarm oranı
    sl = defaultdict(lambda: [0, 0])
    for r, pi in zip(rows, p):
        key = r.get("family") if r["label"] == 1 else r.get("cat")
        if key and key != "none":
            sl[key][0] += int(pi == r["label"])
            sl[key][1] += 1
    if sl:
        m["slices_correct"] = {k: f"{a}/{b}" for k, (a, b) in sorted(sl.items())}
    return m


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model")
    ap.add_argument("--rule", action="store_true")
    ap.add_argument("--name", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--th", type=float, default=0.5)
    ap.add_argument("--max-len", type=int, default=512)
    ap.add_argument("--sets", nargs="*")
    ap.add_argument("--suite", choices=["test", "select"], default="test")
    ap.add_argument("--quiet", action="store_true")
    ap.add_argument("--guard", action="store_true", help="çözücü varyantlarıyla skorla (pinj.guard)")
    ap.add_argument("--ths", type=float, nargs="*", help="ek eşikler (by_threshold)")
    args = ap.parse_args()

    sets = load_sets() if args.suite == "test" else load_selection_sets()
    if args.sets:
        sets = {k: v for k, v in sets.items() if k in args.sets}
    if not args.rule:
        tok = AutoTokenizer.from_pretrained(args.model)
        model = AutoModelForSequenceClassification.from_pretrained(args.model).to(
            "cuda" if torch.cuda.is_available() else "cpu").eval()
    res = {"name": args.name, "model": args.model, "threshold": args.th, "guard": args.guard, "sets": {}}
    for k, rows in sets.items():
        texts = [r["text"] for r in rows]
        t0 = time.time()
        if args.rule:
            sc = rule_scores(texts)
        elif args.guard:
            sc = guarded_scores(model, tok, texts, args.max_len)
        else:
            sc = model_scores(model, tok, texts, args.max_len)
        res["sets"][k] = metrics(rows, sc, args.th)
        if args.ths:  # eşik seçimi için: aynı skorlarla birden çok eşik
            res["sets"][k]["by_threshold"] = {
                str(t): {m: v for m, v in metrics(rows, sc, t).items() if m != "slices_correct"} for t in args.ths}
        r = res["sets"][k]
        if not args.quiet:
          print(f"{args.name:<22} {k:<9} n={r['n']:<5} " + "  ".join(
            f"{m}={r[m]}" for m in ("f1", "precision", "recall", "fpr", "accuracy") if m in r)
              + f"  ({time.time() - t0:.1f}s)")
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    (out / f"{args.name}.json").write_text(json.dumps(res, ensure_ascii=False, indent=1), encoding="utf-8")


if __name__ == "__main__":
    main()