#!/usr/bin/env python3
"""Why did mod-3's jailbreak recall drop, and is mod-2's 99% even real?

Two questions the current benchmark cannot answer, because its jailbreak test
is 200 POSITIVES AND NOTHING ELSE. A model that flags every input scores 100%
on it. So "higher recall" might mean "better" or might mean "flags everything",
and we have no way to tell.

This script fixes both gaps:

1. WHERE DO THE MISSES SIT?  Recall at a fixed 0.5 threshold collapses a
   distribution into one number. If the missed positives score 0.35-0.49 the
   model knows perfectly well and we just picked the wrong operating point,
   which costs nothing to fix. If they score 0.01 it is genuinely blind and
   needs data. Same recall number, completely different problem.

2. JAILBREAK PRECISION.  We score long benign prompts from the SAME forums
   (in-the-wild "regular" configs) as negatives. These are the hard case:
   roleplay-shaped, long, weird, and completely legitimate. Without them you
   are optimising a metric that rewards paranoia.

    python diagnose_jailbreak.py --models mod-2=opus-research/opus-moderation-2 \
                                          mod-3=../moderation3-model
"""
import argparse
import sys
from pathlib import Path

import numpy as np
import torch
from datasets import load_dataset

sys.path.insert(0, str(Path(__file__).resolve().parent))
from benchmark_moderation import load_model  # reuses the arch-aware loader


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--models", nargs="+", required=True, help="NAME=PATH pairs")
    p.add_argument("--n-pos", type=int, default=200)
    p.add_argument("--n-neg", type=int, default=400)
    p.add_argument("--max-len", type=int, default=384)
    p.add_argument("--batch", type=int, default=8,
                   help="small by default: a 1.7B at 384 tokens OOMs an 8GB "
                        "card at the benchmark's batch of 64")
    p.add_argument("--vram-fraction", type=float, default=1.0)
    return p.parse_args()


def main():
    args = parse_args()

    print("loading jailbreak positives...", flush=True)
    jb = load_dataset("jackhhao/jailbreak-classification", split="train")
    pos = [t for t, y in zip(jb["prompt"], jb["type"])
           if str(y).lower().startswith("jail")][-args.n_pos:]

    # The hard negatives: real prompts from the same communities that are NOT
    # jailbreaks. This is the half the benchmark was missing.
    print("loading matched benign negatives (in-the-wild regular)...", flush=True)
    neg = []
    for cfg in ("regular_2023_12_25", "regular_2023_05_07"):
        try:
            ds = load_dataset("TrustAIRLab/in-the-wild-jailbreak-prompts", cfg,
                              split="train")
            neg += [t for t in ds["prompt"] if t and len(t) > 200]
        except Exception as e:
            print(f"  skipped {cfg}: {str(e)[:60]}", flush=True)
    rng = np.random.default_rng(0)
    neg = [neg[i] for i in rng.choice(len(neg), min(len(neg), args.n_neg), replace=False)]
    print(f"  {len(pos)} positives, {len(neg)} long benign negatives\n", flush=True)

    for spec in args.models:
        name, _, path = spec.partition("=")
        print("=" * 66)
        print(f"{name}  ({path})")
        print("=" * 66, flush=True)
        try:
            score, resolve = load_model(path, vram_fraction=args.vram_fraction,
                                        max_len=args.max_len, batch=args.batch)
        except Exception as e:
            print(f"  FAILED to load: {str(e)[:120]}\n", flush=True)
            continue
        if "jailbreaking" not in resolve:
            print("  no jailbreaking label, skipping\n", flush=True)
            continue
        j = resolve["jailbreaking"]

        p_pos = score(pos)[:, j]
        p_neg = score(neg)[:, j]

        # 1. where do the misses sit?
        print("  POSITIVES - score distribution")
        for q in (5, 10, 25, 50, 75, 95):
            print(f"    p{q:<3} {np.percentile(p_pos, q):.3f}")
        missed = p_pos[p_pos < 0.5]
        print(f"  missed at 0.5: {len(missed)}/{len(p_pos)}")
        if len(missed):
            near = (missed >= 0.3).sum()
            print(f"    of those, {near} sit in 0.30-0.49  <- threshold problem, free to fix")
            print(f"              {len(missed)-near} sit below 0.30  <- genuine blindness")
            print(f"    worst miss: {missed.min():.4f}")

        # 2. precision on long benign prompts from the same source
        fp = (p_neg > 0.5).mean()
        print(f"\n  LONG BENIGN NEGATIVES - false positive rate: {fp:.1%}")
        print(f"    p50 {np.percentile(p_neg,50):.3f}   p95 {np.percentile(p_neg,95):.3f}")

        # 3. threshold sweep: recall and FPR move together, show the tradeoff
        print("\n  threshold   recall    FP-rate   F1")
        for t in (0.20, 0.30, 0.40, 0.50, 0.60, 0.70):
            rec = (p_pos > t).mean()
            fpr = (p_neg > t).mean()
            tp, fpc = (p_pos > t).sum(), (p_neg > t).sum()
            prec = tp / (tp + fpc) if tp + fpc else 0.0
            f1 = 2 * prec * rec / (prec + rec) if prec + rec else 0.0
            print(f"    {t:.2f}      {rec:>6.1%}   {fpr:>7.1%}   {f1:.3f}")
        print(flush=True)

        del score
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    print("Read the sweep, not the single number. The best threshold is the one")
    print("where recall is high and the benign FP-rate is still acceptable.")


if __name__ == "__main__":
    main()