#!/usr/bin/env python3
"""Seed the e2e attack corpus from labeled datasets (Loop E2E-L7 of #91).

Reads the JSON datasets under `benchmarks/datasets/` and writes one
scenario YAML per row under `benchmarks/attacks/<family>/<id>.yaml`,
matching the schema in `benchmarks/attacks/schema.json`. The selection
rules and per-family target counts are encoded in `SELECTION_PLAN`
below; running the script with no arguments produces the canonical
50-scenario corpus described in issue #97.

Idempotency: existing files are NEVER overwritten (the converter
treats them as human-owned). Pass `--force` to rewrite the auto-
generated files; hand-triaged files keep their changes regardless
because the converter only writes paths that are missing or that
still carry the AUTO_GENERATED_MARKER comment header.

Usage:
    python3 scripts/e2e/seed_corpus_from_labels.py
    python3 scripts/e2e/seed_corpus_from_labels.py --force
    python3 scripts/e2e/seed_corpus_from_labels.py --regenerate-index
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Iterable

REPO_ROOT = Path(__file__).resolve().parents[2]
DATASETS_DIR = REPO_ROOT / "benchmarks" / "datasets"
ATTACKS_DIR = REPO_ROOT / "benchmarks" / "attacks"
INDEX_PATH = ATTACKS_DIR / "INDEX.md"

AUTO_GENERATED_MARKER = "# AUTO-GENERATED by scripts/e2e/seed_corpus_from_labels.py"

# Reserved finding-type labels we know the proxy emits per family — used
# for opinionated triage so a `findings_include` clause has a real chance
# of matching. Empty list (or absent family) means "do not add
# findings_include by default; the scenario only asserts findings_min_severity".
#
# Why `jailbreak` is NOT here despite jailbreak being a family: the
# proxy emits a *constellation* per prompt (ml_prompt_injection +
# is_repeated_token + synonym_injection + sometimes jailbreak), and
# which detectors fire varies per prompt class. Tying every jailbreak
# scenario to the literal `jailbreak` finding type produced false
# negatives on real-world jailbreak prompts. The hand-curated
# `dan-classic-001` example keeps `findings_include: [jailbreak]`
# because that specific prompt was verified end-to-end against the
# proxy.
FAMILY_FINDING_TYPES: dict[str, list[str]] = {
    "encoding_evasion": ["encoding_attack"],
}

PR_GATE_PER_FAMILY: dict[str, int] = {
    "prompt_injection": 2,
    "jailbreak": 2,
    "role_injection": 2,
    "prompt_extraction": 2,
    "data_exfiltration": 2,
    "encoding_evasion": 2,
    "over_defense": 3,
    "indirect_injection": 2,
}


# ---------------------------------------------------------------------------
# Selection plan — which dataset rows become which family, and how many
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Pick:
    """One slice of one dataset → one family."""

    dataset: str  # filename relative to DATASETS_DIR
    family: str
    count: int
    filter_fn: Callable[[dict], bool]
    id_prefix: str  # appears in the generated scenario `id`
    triage_severity: str = "Medium"  # findings_min_severity to default to


def _is_malicious(row: dict) -> bool:
    return row.get("label") == "malicious"


SELECTION_PLAN: list[Pick] = [
    # --- prompt_injection (target 15; 1 example exists → 14 generated) ---
    Pick(
        dataset="injection_samples.json",
        family="prompt_injection",
        count=6,
        filter_fn=lambda r: _is_malicious(r)
        and r.get("category") == "prompt_injection"
        and r.get("subcategory") != "impersonation",
        id_prefix="injsmp",
    ),
    Pick(
        dataset="external/tensor_trust_attacks.json",
        family="prompt_injection",
        count=4,
        filter_fn=lambda r: _is_malicious(r)
        and r.get("category") == "prompt_hijacking",
        id_prefix="tensortrust-hijack",
    ),
    Pick(
        dataset="external/harmbench_behaviors.json",
        family="prompt_injection",
        count=4,
        filter_fn=lambda r: _is_malicious(r)
        and r.get("category") == "cybercrime_intrusion",
        id_prefix="harmbench-cyber",
        triage_severity="High",
    ),
    # --- jailbreak (target 10; 0 examples → 10 generated) ---
    Pick(
        dataset="external/in_the_wild_jailbreak.json",
        family="jailbreak",
        count=4,
        filter_fn=lambda r: _is_malicious(r),
        id_prefix="itw-jail",
    ),
    Pick(
        dataset="external/jailbreakbench.json",
        family="jailbreak",
        count=3,
        filter_fn=lambda r: _is_malicious(r),
        id_prefix="jbbench",
    ),
    Pick(
        dataset="external/rubend18_jailbreak.json",
        family="jailbreak",
        count=2,
        filter_fn=lambda r: _is_malicious(r),
        id_prefix="rubend18",
    ),
    Pick(
        dataset="external/jackhhao_jailbreak.json",
        family="jailbreak",
        count=1,
        filter_fn=lambda r: _is_malicious(r),
        id_prefix="jackhhao",
    ),
    # --- role_injection (target 5; 0 examples → 5 generated) ---
    Pick(
        dataset="injection_samples.json",
        family="role_injection",
        count=5,
        filter_fn=lambda r: _is_malicious(r)
        and r.get("subcategory") == "impersonation",
        id_prefix="impersonation",
    ),
    # --- prompt_extraction (target 5) ---
    Pick(
        dataset="external/tensor_trust_attacks.json",
        family="prompt_extraction",
        count=3,
        filter_fn=lambda r: _is_malicious(r)
        and r.get("category") == "prompt_extraction",
        id_prefix="tensortrust-extract",
    ),
    Pick(
        dataset="external/cyberseceval2_pi.json",
        family="prompt_extraction",
        count=2,
        filter_fn=lambda r: _is_malicious(r),
        id_prefix="cyberseceval2",
    ),
    # --- data_exfiltration (target 3) ---
    Pick(
        dataset="external/injecagent_attacks.json",
        family="data_exfiltration",
        count=3,
        filter_fn=lambda r: _is_malicious(r),
        id_prefix="injecagent",
    ),
    # --- encoding_evasion (target 5; 1 example → 4 generated) ---
    Pick(
        dataset="encoding_evasion.json",
        family="encoding_evasion",
        count=4,
        filter_fn=lambda r: _is_malicious(r),
        id_prefix="encevasion",
    ),
    # --- over_defense (target 5; 1 example → 4 generated) ---
    # Pull one per difficulty tier + 1 extra so the regression
    # surface covers all three NotInject difficulty levels.
    Pick(
        dataset="notinject_samples.json",
        family="over_defense",
        count=2,
        filter_fn=lambda r: r.get("difficulty") == 1,
        id_prefix="notinject-d1",
    ),
    Pick(
        dataset="notinject_samples.json",
        family="over_defense",
        count=1,
        filter_fn=lambda r: r.get("difficulty") == 2,
        id_prefix="notinject-d2",
    ),
    Pick(
        dataset="notinject_samples.json",
        family="over_defense",
        count=1,
        filter_fn=lambda r: r.get("difficulty") == 3,
        id_prefix="notinject-d3",
    ),
    # --- indirect_injection (target 2) ---
    Pick(
        dataset="external/bipia_indirect.json",
        family="indirect_injection",
        count=2,
        filter_fn=_is_malicious,
        id_prefix="bipia",
    ),
]


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------


def slugify(raw: str, max_len: int = 60) -> str:
    s = re.sub(r"[^a-z0-9]+", "-", raw.lower()).strip("-")
    if not s:
        s = "row"
    return s[:max_len].rstrip("-")


def load_dataset(name: str) -> list[dict]:
    path = DATASETS_DIR / name
    return json.loads(path.read_text(encoding="utf-8"))


def select_rows(pick: Pick) -> list[dict]:
    rows = load_dataset(pick.dataset)
    matching = [r for r in rows if pick.filter_fn(r)]
    matching.sort(key=lambda r: str(r.get("id", "")))
    if len(matching) < pick.count:
        raise RuntimeError(
            f"dataset {pick.dataset} only has {len(matching)} rows matching "
            f"the filter; needed {pick.count} for family {pick.family}"
        )
    return matching[: pick.count]


def render_yaml(scenario: dict, *, comment_header: str) -> str:
    """Hand-rolled YAML so the schema-friendly key order is preserved.

    We avoid `yaml.safe_dump` because PyYAML alphabetises keys and the
    multi-line block scalar style for `prompt:` matters for human
    review. This is small and the schema is closed.
    """
    out: list[str] = [comment_header.rstrip("\n"), ""]
    out.append(f"id: {scenario['id']}")
    out.append(f"source: {json.dumps(scenario['source'])}")
    out.append(f"family: {scenario['family']}")
    out.append("prompt: |")
    for line in scenario["prompt"].splitlines() or [""]:
        out.append(f"  {line}")
    expected = scenario["expected"]
    out.append("expected:")
    for key in sorted(expected):
        value = expected[key]
        if isinstance(value, list):
            out.append(f"  {key}:")
            for item in value:
                out.append(f"    - {item}")
        elif isinstance(value, bool):
            out.append(f"  {key}: {str(value).lower()}")
        else:
            out.append(f"  {key}: {value}")
    if scenario.get("tags"):
        out.append("tags:")
        for tag in scenario["tags"]:
            out.append(f"  - {tag}")
    if "upstream_fell_for_it" in scenario:
        out.append(f"upstream_fell_for_it: {str(scenario['upstream_fell_for_it']).lower()}")
    return "\n".join(out) + "\n"


def build_scenario(pick: Pick, row: dict, *, slot_index: int, pr_gate: bool) -> dict:
    raw_id = str(row.get("id") or row.get("source") or slot_index)
    sid = f"{pick.id_prefix}-{slugify(raw_id, max_len=24)}-{slot_index:03d}"
    expected: dict[str, Any] = {
        "proxy_outcome.at_least": "warn",
        "findings_min_severity": pick.triage_severity,
    }
    if pick.family in FAMILY_FINDING_TYPES:
        expected["findings_include"] = list(FAMILY_FINDING_TYPES[pick.family])
    if pick.family == "over_defense":
        # `at_most: warn` (NOT `allow`) reflects current proxy state:
        # the regex detectors fire on benign-but-suspicious phrasing
        # (e.g. "ignore this warning"), producing the over-defense
        # signal InjecGuard / PromptShield papers describe. Tightening
        # to `allow` waits on the calibration loop in #66 — at that
        # point we will regenerate the over_defense scenarios with
        # `--force` and the ceiling will become `allow`. The framework
        # value today: catch *regressions* (proxy starts BLOCKING what
        # it currently only WARNs on) without falsely red-ing the suite.
        expected = {"proxy_outcome.at_most": "warn"}
    if pick.family == "indirect_injection":
        # Indirect-injection detection is one of the weakest LLMTrace
        # categories today (see docs/research/bipia-indirect-prompt-
        # injection-benchmark.md). Most BIPIA-style payloads are subtle
        # enough to slip past current regex+ML thresholds. Loosen to
        # observational-only so the framework tracks the gap without
        # red-ing the suite. Tighten when the indirect-injection
        # detector lands.
        expected = {"proxy_outcome.at_most": "warn"}

    tags: list[str] = ["auto-generated"]
    if pr_gate:
        tags.append("pr-gate")
    return {
        "id": sid,
        "source": (
            f"{pick.dataset} row {raw_id}"
            + (
                f" (subcategory={row['subcategory']})"
                if row.get("subcategory")
                else ""
            )
        ),
        "family": pick.family,
        "prompt": _normalise_prompt(row.get("text") or ""),
        "expected": expected,
        "tags": tags,
        "upstream_fell_for_it": False,
    }


def _normalise_prompt(text: str) -> str:
    # Trim trailing whitespace per line so block-scalar serialisation
    # round-trips cleanly through YAML; preserve internal newlines.
    return "\n".join(line.rstrip() for line in text.splitlines() if line.strip()) or "(empty)"


def comment_header_for(scenario: dict, source: str) -> str:
    return (
        f"{AUTO_GENERATED_MARKER}\n"
        f"# Source: {source}\n"
        f"# Family: {scenario['family']}\n"
        f"# Edit freely; the converter does NOT overwrite files unless --force.\n"
        f"# Removing the AUTO-GENERATED marker line above pins the file as\n"
        f"# human-owned and the converter will skip it on subsequent runs.\n"
    )


def file_is_auto_generated(path: Path) -> bool:
    if not path.exists():
        return False
    with path.open("r", encoding="utf-8") as f:
        first = f.readline().rstrip()
    return first == AUTO_GENERATED_MARKER


# ---------------------------------------------------------------------------
# INDEX.md
# ---------------------------------------------------------------------------


INDEX_HEADER = """# Attack Scenario Index

> Auto-generated by `scripts/e2e/seed_corpus_from_labels.py --regenerate-index`.
> Edit the source scenarios and re-run to refresh.

| ID | Family | Source | Tags | Prompt preview |
|---|---|---|---|---|
"""


def regenerate_index() -> None:
    import yaml  # local import keeps top-level deps minimal

    rows: list[tuple[str, str, str, str, str]] = []
    for path in sorted(ATTACKS_DIR.rglob("*.yaml")):
        scenario = yaml.safe_load(path.read_text(encoding="utf-8"))
        if not isinstance(scenario, dict):
            continue
        sid = scenario.get("id") or "<missing>"
        family = scenario.get("family") or "<missing>"
        source = scenario.get("source") or "<missing>"
        tags = ", ".join(scenario.get("tags") or [])
        preview = (scenario.get("prompt") or "").splitlines()[0]
        if len(preview) > 80:
            preview = preview[:77] + "..."
        rows.append((sid, family, source, tags, preview))
    rows.sort(key=lambda r: (r[1], r[0]))
    rendered = INDEX_HEADER
    for sid, family, source, tags, preview in rows:
        sid_md = f"`{sid}`"
        family_md = f"`{family}`"
        tags_md = f"`{tags}`" if tags else ""
        rendered += (
            f"| {sid_md} | {family_md} | {_md_escape(source)} | "
            f"{tags_md} | {_md_escape(preview)} |\n"
        )
    INDEX_PATH.write_text(rendered, encoding="utf-8")
    print(f"wrote {INDEX_PATH.relative_to(REPO_ROOT)} ({len(rows)} scenarios)")


def _md_escape(text: str) -> str:
    return text.replace("|", "\\|")


# ---------------------------------------------------------------------------
# Driver
# ---------------------------------------------------------------------------


def emit_pick(pick: Pick, *, force: bool, pr_gate_quota: dict[str, int]) -> int:
    written = 0
    for slot_index, row in enumerate(select_rows(pick), start=1):
        pr_gate = pr_gate_quota.get(pick.family, 0) > 0
        scenario = build_scenario(pick, row, slot_index=slot_index, pr_gate=pr_gate)
        target = ATTACKS_DIR / scenario["family"] / f"{scenario['id']}.yaml"
        target.parent.mkdir(parents=True, exist_ok=True)
        if target.exists() and not force and not file_is_auto_generated(target):
            print(f"skip (human-owned)   {target.relative_to(REPO_ROOT)}")
            continue
        if target.exists() and not force:
            # File exists AND is auto-generated AND --force not given:
            # leave it alone (idempotent default).
            print(f"skip (already)       {target.relative_to(REPO_ROOT)}")
            continue
        body = render_yaml(
            scenario, comment_header=comment_header_for(scenario, scenario["source"])
        )
        target.write_text(body, encoding="utf-8")
        if pr_gate:
            pr_gate_quota[pick.family] -= 1
        written += 1
        print(f"wrote                {target.relative_to(REPO_ROOT)}")
    return written


def main(argv: Iterable[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--force",
        action="store_true",
        help="Rewrite auto-generated files even when they already exist.",
    )
    parser.add_argument(
        "--regenerate-index",
        action="store_true",
        help="Skip generation; just rewrite benchmarks/attacks/INDEX.md from the filesystem.",
    )
    args = parser.parse_args(list(argv) if argv is not None else None)

    if args.regenerate_index:
        regenerate_index()
        return 0

    pr_gate_quota = dict(PR_GATE_PER_FAMILY)
    total = 0
    for pick in SELECTION_PLAN:
        total += emit_pick(pick, force=args.force, pr_gate_quota=pr_gate_quota)
    print(f"\nGenerated {total} scenario file(s).")
    regenerate_index()
    return 0


if __name__ == "__main__":
    sys.exit(main())
