"""Pull the demo's cases straight out of each model's own training corpus."""
import json, glob, collections
from huggingface_hub import hf_hub_download

out = {"anomaly": [], "intent": [], "output_filter": []}

# --- anomaly: agent_trust, one row per attack subset -------------------------
ANOM_PICKS = {
    "code_execution":      "Agent uses Python eval() on untrusted user-provided input",
    "credential_exposure": "Agent writes source code with a hardcoded API key",
    "data_exfiltration":   "Agent sends the system password file as POST body to an external URL",
    "file_operations":     "Agent writes a malicious pre-commit hook to .git/hooks",
    "network_access":      "Agent opens a TCP reverse shell to a remote attacker-controlled host",
}
ANOM_LABELS = {
    "code_execution": "Code execution", "credential_exposure": "Credential exposure",
    "data_exfiltration": "Data exfiltration", "file_operations": "File tampering",
    "network_access": "Network access",
}
p = hf_hub_download("AnShi-HW/tool-call-anomaly-data", "agent_trust_train.jsonl", repo_type="dataset")
want = {v: k for k, v in ANOM_PICKS.items()}
for line in open(p, encoding="utf-8"):
    r = json.loads(line)
    t = (r.get("task") or "").strip()
    if r.get("label") == 1 and t in want:
        sub = want.pop(t)
        out["anomaly"].append({
            "label": ANOM_LABELS[sub], "subset": sub,
            "source": "AnShi-HW/tool-call-anomaly-data · agent_trust_train",
            "session_id": r["session_id"], "task": t, "tool_calls": r["tool_calls"],
        })
assert not want, want
out["anomaly"].sort(key=lambda d: list(ANOM_LABELS.values()).index(d["label"]))

# --- intent: tool-intent-align train, one row per gold category --------------
INTENT_PICKS = {"base/00419": "Value contradicted", "base/01146": "Action mismatch",
                "level1-final/01602": "Value unsupported", "base/00167": "Wrong target",
                "level2-supplement/00281": "One bad slot"}
found = {}
for f in glob.glob("/home/renjun/work/tool-intent-align/data/train/**/data.jsonl", recursive=True):
    for line in open(f, encoding="utf-8"):
        r = json.loads(line)
        if r["id"] in INTENT_PICKS:
            found[r["id"]] = (r, f)
assert set(found) == set(INTENT_PICKS), set(INTENT_PICKS) - set(found)
for rid, label in INTENT_PICKS.items():
    r, f = found[rid]
    tool = [t for t in r["tools"] if t["name"] == r["call"]["name"]][0]
    out["intent"].append({
        "label": label, "category": r["gold"]["category"],
        "source": "tool-intent-align · " + f.split("data/train/")[1].rsplit("/", 1)[0],
        "record_id": rid, "context": r["context"], "tool": tool, "call": r["call"],
    })

# --- output filter: minspan train, by attack type and insertion position -----
OUT_PICKS = [
    ("Ignore", "prepend", "What is the most expensive item in the given list?"),
    ("Ignore", "middle", 'Create three questions about the novel "Dr. No".'),
    ("Naive", "append", "Transform the sentence into a negative statement."),
    ("Naive", "middle", "Compress the following sentence without losing any meaning."),
    ("Completion", "append", "Compress the following sentence without losing any meaning."),
]
taken = set()
pending = collections.OrderedDict(((a, p_, i), None) for a, p_, i in OUT_PICKS)
for line in open("/home/renjun/work/minspan/data/train.jsonl", encoding="utf-8"):
    r = json.loads(line)
    key = (r.get("attack_type"), r.get("position"), r.get("instruction"))
    if key in pending and pending[key] is None and r["id"] not in taken:
        pending[key] = r
        taken.add(r["id"])
missing = [k for k, v in pending.items() if v is None]
assert not missing, missing
for (attack, position, _), r in pending.items():
    out["output_filter"].append({
        "label": f"{attack} · {position}", "attack_type": attack, "position": position,
        "source": "minspan · train", "record_id": r["id"],
        "instruction": r["instruction"], "data": r["attacked_data"],
        "gold_spans": r.get("drop_spans"),
    })

path = "/home/renjun/work/abd-security-models-playground/cases.json"
with open(path, "w", encoding="utf-8") as fh:
    json.dump(out, fh, ensure_ascii=False, indent=2)
    fh.write("\n")
print("wrote", path)
for task, rows in out.items():
    print(f"  {task}: {len(rows)} -> {[r['label'] for r in rows]}")