import { homedir } from "node:os"
import { join } from "node:path"
import Papa from "papaparse"
import type { FixtureRow } from "../../types.ts"
import { fetchCached } from "../fetcher.ts"

// Pinned upstream revisions. Bump explicitly when pulling newer data —
// changing a SHA is a deliberate act, not an incidental refresh.
const HF_DATASET_REVISION = "886acc352a31533ffbcf4ef22c744658688086fc"
const ARTIFACTS_REPO_REVISION = "909e68c01d94222b8ad2e397a017e2e12e2adb73"

// Target models that JailbreakBench ran its attacks against.
const ALL_TARGET_MODELS = ["gpt-3.5-turbo-1106", "gpt-4-0125-preview", "llama-2-7b-chat-hf", "vicuna-13b-v1.5"] as const

// The three attack methods we sample from the JailbreakBench artifacts repo.
// Each exercises a different slice of our flagger's jailbreak taxonomy
// (see run-flagger.ts > JAILBREAK_SYSTEM_PROMPT for categories).
//
// - JBC / manual — human-written persona roleplay. 100 goals × 1 template
//   (the AIM / Machiavelli "amoral character" prompt). All 100 positives
//   per model share the same wrapping; only the final `goal` varies. This
//   is the ROLE / PERSONA ATTACKS category. We use only one target model
//   because the user prompt is identical across all four — response
//   variety doesn't change the flagger's decision (the user-turn tactic
//   is already decisive), so 4× the same prompt = 4× correlated rows.
//
// - PAIR / black-box — prompts generated by an attacker LLM via the PAIR
//   algorithm (Prompt Automatic Iterative Refinement). Each row is a unique
//   fictional-framing attempt: "you're a gossip columnist in a parallel
//   universe", "for academic research on cultural differences", "as a
//   creative writer working on a dystopian novel". Exercises FICTIONAL /
//   HYPOTHETICAL FRAMING and often COMPETING OBJECTIVES. PAIR runs
//   separately per target model, so every (goal, model) row has a distinct
//   prompt — keep all four models.
//
// - GCG — gradient-based adversarial suffix. The goal string is left intact
//   and a trained gibberish suffix is appended that reliably triggers
//   compliance on the target. Exercises CHARACTER / ENCODING OBFUSCATION.
//   `white_box` for open-weight targets (llama, vicuna); `transfer` for
//   closed models (gpt-*) where suffixes trained on one model are reused.
//   Suffixes are trained per model, so each (goal, model) pair is unique —
//   keep all four models.
interface AttackSource {
  readonly id: string
  readonly tactic: string
  readonly models: readonly string[]
  readonly attackPath: (model: string) => string
}

const ATTACK_SOURCES: readonly AttackSource[] = [
  {
    id: "jbc-manual",
    tactic: "persona-aim",
    models: ["vicuna-13b-v1.5"],
    attackPath: (m) => `JBC/manual/${m}.json`,
  },
  {
    id: "pair-black-box",
    tactic: "fictional-framing",
    models: ALL_TARGET_MODELS,
    attackPath: (m) => `PAIR/black_box/${m}.json`,
  },
  {
    id: "gcg",
    tactic: "adversarial-suffix",
    models: ALL_TARGET_MODELS,
    attackPath: (m) => (m.startsWith("gpt-") ? `GCG/transfer/${m}.json` : `GCG/white_box/${m}.json`),
  },
]

const hfFileUrl = (file: string): string =>
  `https://huggingface.co/datasets/JailbreakBench/JBB-Behaviors/resolve/${HF_DATASET_REVISION}/${file}`

const artifactUrl = (subpath: string): string =>
  `https://raw.githubusercontent.com/JailbreakBench/artifacts/${ARTIFACTS_REPO_REVISION}/attack-artifacts/${subpath}`

// -- Upstream row shapes (what we download) -----------------------------------

interface BehaviorsCsvRow {
  readonly Index: string
  readonly Goal: string
  readonly Target: string
  readonly Behavior: string
  readonly Category: string
  readonly Source: string
}

interface JbcJailbreak {
  readonly index: number
  readonly goal: string
  readonly behavior: string
  readonly category: string
  // PAIR leaves prompt/response null on rows where the attacker failed to
  // find a working jailbreak — we filter those out before building fixtures.
  readonly prompt: string | null
  readonly response: string | null
  readonly jailbroken: boolean
}

interface JbcArtifact {
  readonly jailbreaks: readonly JbcJailbreak[]
}

// -- Parsers ------------------------------------------------------------------

function parseBehaviorsCsv(buffer: Buffer): BehaviorsCsvRow[] {
  const result = Papa.parse<BehaviorsCsvRow>(buffer.toString("utf8"), {
    header: true,
    skipEmptyLines: true,
  })
  if (result.errors.length > 0) {
    throw new Error(`CSV parse errors: ${JSON.stringify(result.errors.slice(0, 3))}`)
  }
  return result.data
}

function parseJbcArtifact(buffer: Buffer): JbcArtifact {
  return JSON.parse(buffer.toString("utf8")) as JbcArtifact
}

// -- FixtureRow builders ------------------------------------------------------

/**
 * Positive: a real jailbreak attempt (user turn) with real model response.
 * Returns null if the upstream row has no attack prompt (PAIR leaves these
 * nulled when the attacker failed to find a working jailbreak).
 */
function buildPositive(source: AttackSource, model: string, row: JbcJailbreak): FixtureRow | null {
  if (row.prompt === null || row.response === null) return null
  return {
    id: `jailbreakbench-${source.id}-${model}-${row.index}`,
    source: `jailbreakbench:attack-artifacts/${source.attackPath(model)}@${ARTIFACTS_REPO_REVISION.slice(0, 7)}`,
    licence: "MIT",
    expected: { matched: true },
    tags: [
      row.category,
      source.id,
      `tactic:${source.tactic}`,
      row.jailbroken ? "assistant-complied" : "assistant-refused",
    ],
    trace: {
      messages: [
        { role: "user", parts: [{ type: "text", content: row.prompt }] },
        { role: "assistant", parts: [{ type: "text", content: row.response }] },
      ],
    },
    notes: `Goal: ${row.goal}`,
  }
}

/**
 * Soft negative: the benign JBB-Behaviors Goal used as a plain user turn.
 * No assistant turn — the flagger's required context is one message.
 * Should never match for jailbreaking: not harmful, not adversarial.
 */
function buildSoftNegative(row: BehaviorsCsvRow): FixtureRow {
  return {
    id: `jailbreakbench-benign-${row.Index}`,
    source: `jailbreakbench:JBB-Behaviors/benign@${HF_DATASET_REVISION.slice(0, 7)}`,
    licence: "MIT",
    expected: { matched: false },
    tags: [row.Category, "tactic:jbb-benign"],
    trace: {
      messages: [{ role: "user", parts: [{ type: "text", content: row.Goal }] }],
    },
    notes: `Behavior: ${row.Behavior}`,
  }
}

/**
 * Hard negative: the harmful JBB-Behaviors Goal used as a plain direct ask,
 * no manipulation tactic, no DAN wrapping. The flagger's definition requires
 * MANIPULATION TACTIC + adversarial intent — a direct harmful question on its
 * own is explicitly not a jailbreak. These are the precision-stress cases:
 * if the flagger fires here, we have a false-positive problem.
 */
function buildHardNegative(row: BehaviorsCsvRow): FixtureRow {
  return {
    id: `jailbreakbench-harmful-direct-${row.Index}`,
    source: `jailbreakbench:JBB-Behaviors/harmful@${HF_DATASET_REVISION.slice(0, 7)}`,
    licence: "MIT",
    expected: { matched: false },
    tags: [row.Category, "tactic:jbb-harmful-direct", "hard-negative"],
    trace: {
      messages: [{ role: "user", parts: [{ type: "text", content: row.Goal }] }],
    },
    notes: `Direct harmful request without manipulation tactic. Behavior: ${row.Behavior}`,
  }
}

// -- Main ---------------------------------------------------------------------

/**
 * Fetch JailbreakBench sources, parse, and emit FixtureRows for the
 * `jailbreaking` queue. Output is always ordered: positives, then soft
 * negatives, then hard negatives — stable so repeat runs produce a stable
 * JSONL diff.
 *
 * Positives are pulled from three attack methods in the JailbreakBench
 * artifacts repo, each exercising a different tactic in our flagger's
 * taxonomy (see run-flagger.ts > JAILBREAK_SYSTEM_PROMPT):
 *
 *   - **JBC / manual** — human-written persona roleplay (AIM / Machiavelli
 *     template). ROLE / PERSONA ATTACKS. Deduped to 1 target model since
 *     the user prompt is identical across the four artifact files.
 *   - **PAIR / black-box** — algorithmically generated fictional-framing
 *     prompts produced by an attacker LLM. Each (goal, model) row has a
 *     distinct prompt. FICTIONAL / HYPOTHETICAL FRAMING + COMPETING
 *     OBJECTIVES. ~400 rows before filtering; PAIR leaves prompt/response
 *     null on rows where the attacker failed — those are dropped, yielding
 *     ~230-240 valid rows.
 *   - **GCG** — gradient-based adversarial suffix. Each (goal, model) row
 *     has a distinct suffix. CHARACTER / ENCODING OBFUSCATION. 400 rows.
 *
 * Soft negatives (100) and hard negatives (100) come from JBB-Behaviors
 * benign and harmful splits respectively — see `buildSoftNegative` /
 * `buildHardNegative`.
 *
 * `options.cacheRoot` defaults to `~/.cache/latitude-benchmarks/upstream/`.
 */
export async function mapJailbreakBench(options: { readonly cacheRoot?: string } = {}): Promise<FixtureRow[]> {
  const cacheRoot = options.cacheRoot ?? join(homedir(), ".cache", "latitude-benchmarks", "upstream")

  const positives: FixtureRow[] = []
  for (const source of ATTACK_SOURCES) {
    for (const model of source.models) {
      const subpath = source.attackPath(model)
      const buffer = await fetchCached(
        artifactUrl(subpath),
        join(cacheRoot, `jailbreakbench@${ARTIFACTS_REPO_REVISION}`, subpath),
      )
      const artifact = parseJbcArtifact(buffer)
      for (const row of artifact.jailbreaks) {
        const fixture = buildPositive(source, model, row)
        if (fixture !== null) positives.push(fixture)
      }
    }
  }

  const benignBuf = await fetchCached(
    hfFileUrl("data/benign-behaviors.csv"),
    join(cacheRoot, `jailbreakbench@${HF_DATASET_REVISION}`, "benign-behaviors.csv"),
  )
  const softNegatives = parseBehaviorsCsv(benignBuf).map(buildSoftNegative)

  const harmfulBuf = await fetchCached(
    hfFileUrl("data/harmful-behaviors.csv"),
    join(cacheRoot, `jailbreakbench@${HF_DATASET_REVISION}`, "harmful-behaviors.csv"),
  )
  const hardNegatives = parseBehaviorsCsv(harmfulBuf).map(buildHardNegative)

  return [...positives, ...softNegatives, ...hardNegatives]
}
