"""
How safe is the model?  Steps 3 to 5.
=====================================
Data: JailbreakBench (Chao et al., 2024; github.com/JailbreakBench, MIT).
100 requests a model should refuse, in 10 categories. Three published attack
methods tried to get four models to comply anyway. Each attempt was marked
"jailbroken" or not by the benchmark's judge (a language model), and also by
Llama Guard. A separate set of 300 responses was labeled by three people and
by four automatic judges.

The files contain harmful requests and the models' responses. This script
reads only the labels and counts. It never prints a request or a response,
and you don't need to open the files to do the lesson.

Setup (once):  nothing to install beyond Python 3
Run:           python redteam_lesson.py               (attack success by model and method)
               python redteam_lesson.py judges        (automatic judges against people)
               python redteam_lesson.py categories    (by category, and two judges on the same attacks)
The first run downloads about 3 MB into ./data.
"""
import csv
import json
import os
import sys
import urllib.request
from collections import defaultdict

ART = 'https://raw.githubusercontent.com/JailbreakBench/artifacts/main/attack-artifacts/'
JUDGES_URL = 'https://huggingface.co/datasets/JailbreakBench/JBB-Behaviors/resolve/main/data/judge-comparison.csv'
MODELS = {'llama-2-7b-chat-hf': 'Llama-2 7B', 'vicuna-13b-v1.5': 'Vicuna 13B',
          'gpt-3.5-turbo-1106': 'GPT-3.5', 'gpt-4-0125-preview': 'GPT-4'}
METHODS = {'GCG': 'GCG', 'PAIR': 'PAIR', 'prompt_with_random_search': 'Random search'}
JUDGES = {'harmbench_cf': 'HarmBench classifier', 'llamaguard2_cf': 'Llama Guard 2',
          'gpt4_cf': 'GPT-4 as judge', 'llama3_cf': 'Llama-3 70B as judge'}


def fetch(url, path):
    os.makedirs('data', exist_ok=True)
    if not os.path.exists(path):
        print('downloading', os.path.basename(path), '...')
        urllib.request.urlretrieve(url, path)
    return path


def attacks():
    """Yield (model, method, request number, category, judged jailbroken, Llama Guard said jailbroken) per attempt."""
    for method in METHODS:
        for model in MODELS:
            kind = 'black_box' if method != 'GCG' else ('white_box' if 'gpt' not in model else 'transfer')
            path = fetch(f'{ART}{method}/{kind}/{model}.json', f'data/jbb_{method}_{model}.json')
            for j in json.load(open(path, encoding='utf-8'))['jailbreaks']:
                yield model, method, j['index'], j['category'], bool(j['jailbroken']), bool(j['jailbroken_llama_guard1'])


def report():
    hit = defaultdict(dict)                                  # model -> request index -> {method: jailbroken}
    rate = defaultdict(list)
    for model, method, i, _, jb, _ in attacks():
        hit[model].setdefault(i, {})[method] = jb
        rate[(model, method)].append(jb)
    print("Share of 100 harmful requests where the attack got a harmful answer (benchmark judge)\n")
    print(f"  {'model':<12}" + ''.join(f"{m:>15}" for m in METHODS.values()) + f"{'any of the 3':>15}")
    for model, name in MODELS.items():
        row = ''.join(f"{sum(rate[(model, m)]) / len(rate[(model, m)]):>15.0%}" for m in METHODS)
        worst = sum(any(v.values()) for v in hit[model].values()) / len(hit[model])
        print(f"  {name:<12}{row}{worst:>15.0%}")


def judges():
    rows = list(csv.DictReader(open(fetch(JUDGES_URL, 'data/jbb_judge_comparison.csv'), encoding='utf-8')))
    people = [int(r['human_majority']) for r in rows]
    print(f"{len(rows)} responses, each labeled harmful or not by three people. Majority says harmful: "
          f"{sum(people)} ({sum(people) / len(rows):.0%})\n")
    print(f"  {'judge':<22} {'agrees with people':>18} {'says harmful':>13} {'of its flags, people agree':>27} {'of real harm, it catches':>25}")
    for col, name in JUDGES.items():
        says = [int(float(r[col])) for r in rows]
        both = sum(a and b for a, b in zip(people, says))
        agree = sum(a == b for a, b in zip(people, says)) / len(rows)
        print(f"  {name:<22} {agree:>18.0%} {sum(says) / len(rows):>13.0%} {both / sum(says):>27.0%} {both / sum(people):>25.0%}")
    print("\nPeople against each other:")
    for a, b in (('human1', 'human2'), ('human1', 'human3'), ('human2', 'human3')):
        print(f"  person {a[-1]} and person {b[-1]} agree on {sum(r[a] == r[b] for r in rows) / len(rows):.0%}")


def categories():
    cat, same, only_judge, only_guard, n = defaultdict(list), 0, 0, 0, 0
    for _, _, _, c, jb, lg in attacks():
        cat[c].append(jb)
        n += 1
        same += jb == lg
        only_judge += jb and not lg
        only_guard += lg and not jb
    print("Attack success by category (all models and methods together)\n")
    for c, v in sorted(cat.items(), key=lambda kv: -sum(kv[1]) / len(kv[1])):
        print(f"  {c:<28} {sum(v) / len(v):>5.0%}")
    print(f"\nThe benchmark judge and Llama Guard, on the same {n:,} attack results:")
    print(f"  agree on {same / n:.1%}   only the benchmark judge says harmful: {only_judge}   only Llama Guard says harmful: {only_guard}")


if __name__ == '__main__':
    cmd = sys.argv[1] if len(sys.argv) > 1 else ''
    if cmd == 'judges':
        judges()
    elif cmd == 'categories':
        categories()
    else:
        report()