3-generate_distractor_facts.py

#!/usr/bin/env python3
"""Generate v24 distractor conversations for noisy longitudinal evaluation."""

from __future__ import annotations

import argparse
import json
import os
import re
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Callable, Iterable, Sequence

from tqdm import tqdm


BASE_DIR = Path(__file__).resolve().parent
DEFAULT_EVIDENCE_DIR = BASE_DIR / "generated_result" / "2-evidence_conversation"
DEFAULT_OUTPUT_DIR = BASE_DIR / "generated_result" / "3-distractor_facts"
GENERATION_TEMPLATE_PATH = BASE_DIR / "template" / "3-generate_distractor_segment.md"
VERIFY_TEMPLATE_PATH = BASE_DIR / "template" / "3.5-llm_verify_distractor_segment.md"
DEFAULT_GENERATION_MODEL = os.getenv("DISTRACTOR_GENERATION_MODEL", "claude-opus-4-6")
DEFAULT_VERIFY_MODEL = os.getenv("DISTRACTOR_VERIFY_MODEL", "gpt-4o")
DEFAULT_MAX_ATTEMPTS = int(os.getenv("DISTRACTOR_MAX_ATTEMPTS", "3"))
DEFAULT_MAX_PER_SAMPLE = int(os.getenv("DISTRACTOR_MAX_PER_SAMPLE", "6"))
DEFAULT_CONCURRENCY = int(os.getenv("DISTRACTOR_CONCURRENCY", "4"))
DEFAULT_REQUEST_TIMEOUT = float(os.getenv("DISTRACTOR_REQUEST_TIMEOUT", "180"))
DEFAULT_REQUEST_RETRIES = int(os.getenv("DISTRACTOR_REQUEST_RETRIES", "3"))
DEFAULT_REQUEST_RETRY_SLEEP_SECONDS = float(
    os.getenv("DISTRACTOR_REQUEST_RETRY_SLEEP_SECONDS", "2")
)
DEFAULT_MAX_TOKENS = int(os.getenv("DISTRACTOR_MAX_TOKENS", "8000"))
DEFAULT_TEMPERATURE = float(os.getenv("DISTRACTOR_TEMPERATURE", "0.2"))
USE_LLM_BY_DEFAULT = os.getenv("DISTRACTOR_USE_LLM", "1").strip().lower() not in {
    "0",
    "false",
    "no",
    "off",
}
BASE_URL = os.getenv("LLM_BASE_URL", "").rstrip("/")
KEY_FILE_PATH = Path(
    os.getenv("DISTRACTOR_KEY_FILE", os.getenv("EVIDENCE_KEY_FILE", str(BASE_DIR / "api.md")))
)
API_KEY_FALLBACK_PATH = BASE_DIR / "key.md"


LlmCaller = Callable[[str, str], str]
ProgressFactory = Callable[..., Iterable[Any]]


DISTRACTOR_TYPES = {
    "third_party_same_target",
    "near_miss_third_party",
    "hypothetical_self",
    "assistant_echo_wrong_value",
    "stale_neighbor",
    "retained_neighbor",
    "tentative_trap_neighbor",
}
ASSISTANT_ONLY_TYPES = {"assistant_echo_wrong_value"}
STATE_NEUTRAL_MARKERS = (
    "third party",
    "coworker",
    "neighbor",
    "sibling",
    "friend",
    "assistant-only",
    "assistant only",
    "not the benchmark user",
    "not the user",
    "hypothetical",
    "not adopted",
    "not confirmed",
    "retained",
    "someone else",
    "colleague",
)

# Appended to retry feedback so the generator clearly understands the constraints
# that most often cause verifier rejections, instead of only seeing terse issue
# strings. State-neutrality and span validity are now judged by the LLM verifier.
RETRY_CONSTRAINTS_REMINDER = (
    "Constraints to satisfy on the next attempt:\n"
    "- Copy each distractor_span.quote VERBATIM (character-for-character) from one "
    "message inside that distractor's own dialogue. Do not paraphrase, rewrite "
    "pronouns (e.g. 'he' -> a name), reorder words, or change capitalization.\n"
    "- The distractor_value must appear exactly, as written, inside its "
    "distractor_span.quote.\n"
    "- Scope every distractor value to a third party, an assistant-only draft, or a "
    "hypothetical; never confirm or restore it as the benchmark user's own value, "
    "and state this clearly in state_neutrality_rationale.\n"
    "- Never reuse the user's real (forbidden) or forgotten gold values as a "
    "distractor_value.\n"
    "- Each dialogue must have 2-6 messages, start with the user, end with the "
    "assistant, and strictly alternate roles."
)


def read_text(path: Path) -> str:
    return path.read_text(encoding="utf-8")


def read_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def is_final_evidence_file(path: Path) -> bool:
    return (
        path.is_file()
        and path.suffix == ".json"
        and "_attempt_" not in path.name
        and "_llm_verify" not in path.name
        and not path.name.startswith("question_answer_translation_cache")
        and not path.name.startswith("report_translation_cache")
    )


def iter_evidence_files(evidence_dir: Path) -> list[Path]:
    if not evidence_dir.is_dir():
        raise FileNotFoundError(f"Evidence directory not found: {evidence_dir}")
    return sorted(path for path in evidence_dir.glob("*.json") if is_final_evidence_file(path))


def parse_api_key_text(text: str) -> str | None:
    for raw_line in text.splitlines():
        stripped = raw_line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        if stripped.startswith("export "):
            stripped = stripped.removeprefix("export ").strip()
        if "=" in stripped:
            stripped = stripped.split("=", 1)[1].strip()
        key = stripped.strip().strip('"').strip("'")
        if key:
            return key
    return None


def read_api_key_file(path: Path) -> str | None:
    if not path.exists():
        return None
    return parse_api_key_text(path.read_text(encoding="utf-8"))


def get_api_key() -> str | None:
    return (
        os.getenv("MEMTENSOR_API_KEY")
        or os.getenv("OPENAI_API_KEY")
        or os.getenv("ANTHROPIC_API_KEY")
        or os.getenv("API_KEY")
        or read_api_key_file(KEY_FILE_PATH)
        or read_api_key_file(API_KEY_FALLBACK_PATH)
    )


def parse_api_keys_text(text: str) -> list[str]:
    """Return every key in a key file (one per non-comment line)."""
    keys: list[str] = []
    for raw_line in text.splitlines():
        stripped = raw_line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        if stripped.startswith("export "):
            stripped = stripped.removeprefix("export ").strip()
        if "=" in stripped:
            stripped = stripped.split("=", 1)[1].strip()
        key = stripped.strip().strip('"').strip("'")
        if key:
            keys.append(key)
    return keys


def get_api_keys() -> list[str]:
    """All candidate keys, in priority order, de-duplicated.

    Some keys on the shared gateway can be individually rate-limited / 503'd
    while others work, so call_llm rotates across this list before giving up.
    """
    candidates: list[str] = []
    env_key = (
        os.getenv("MEMTENSOR_API_KEY")
        or os.getenv("OPENAI_API_KEY")
        or os.getenv("ANTHROPIC_API_KEY")
        or os.getenv("API_KEY")
    )
    if env_key:
        candidates.append(env_key.strip())
    for path in (KEY_FILE_PATH, API_KEY_FALLBACK_PATH):
        if path.exists():
            candidates.extend(parse_api_keys_text(path.read_text(encoding="utf-8")))
    seen: set[str] = set()
    unique: list[str] = []
    for key in candidates:
        if key and key not in seen:
            seen.add(key)
            unique.append(key)
    return unique


def chat_completions_url(base_url: str) -> str:
    if base_url.endswith("/v1"):
        return f"{base_url}/chat/completions"
    return f"{base_url}/v1/chat/completions"


def call_llm(model: str, prompt: str) -> str:
    if not BASE_URL:
        raise RuntimeError(
            "Missing LLM_BASE_URL. Set it to an OpenAI-compatible chat-completions endpoint (e.g. https://api.openai.com or your own gateway/router) before running this script."
        )
    api_keys = get_api_keys()
    if not api_keys:
        raise RuntimeError(
            "Missing API key for distractor generation. Set OPENAI_API_KEY/"
            "MEMTENSOR_API_KEY/API_KEY or create api.md (key.md also supported)."
        )
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": DEFAULT_TEMPERATURE,
        "max_tokens": DEFAULT_MAX_TOKENS,
    }
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    last_error: Exception | None = None
    # Each round tries every key; a key that is individually rate-limited / 503'd
    # is skipped in favour of the next, so one bad key no longer blocks the run.
    for attempt in range(1, DEFAULT_REQUEST_RETRIES + 1):
        for api_key in api_keys:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            }
            try:
                request = urllib.request.Request(
                    chat_completions_url(BASE_URL),
                    data=data,
                    headers=headers,
                    method="POST",
                )
                with urllib.request.urlopen(request, timeout=DEFAULT_REQUEST_TIMEOUT) as response:
                    response_payload = json.loads(response.read().decode("utf-8"))
                return response_payload["choices"][0]["message"]["content"]
            except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError) as exc:
                last_error = exc
                continue
        if attempt < DEFAULT_REQUEST_RETRIES:
            time.sleep(DEFAULT_REQUEST_RETRY_SLEEP_SECONDS * attempt)
    raise RuntimeError(
        f"LLM request failed after {DEFAULT_REQUEST_RETRIES} round(s) over "
        f"{len(api_keys)} key(s): {last_error}"
    )


def parse_json_object(text: str) -> dict[str, Any] | None:
    stripped = text.strip()
    try:
        parsed = json.loads(stripped)
        if isinstance(parsed, dict):
            return parsed
    except json.JSONDecodeError:
        pass
    start = stripped.find("{")
    end = stripped.rfind("}")
    if start == -1 or end == -1 or end <= start:
        return None
    try:
        parsed = json.loads(stripped[start : end + 1])
    except json.JSONDecodeError:
        return None
    return parsed if isinstance(parsed, dict) else None


def normalize_text(value: Any) -> str:
    text = str(value).strip().lower()
    text = re.sub(r"\s+", " ", text)
    text = re.sub(r"[\"'`.,;:!?，。；：！？]+$", "", text)
    return text


def value_strings(value: Any) -> list[str]:
    if value is None:
        return []
    if isinstance(value, str):
        stripped = value.strip()
        return [stripped] if stripped else []
    if isinstance(value, (int, float)):
        return [str(value)]
    if isinstance(value, list):
        strings: list[str] = []
        for item in value:
            strings.extend(value_strings(item))
        return strings
    if isinstance(value, dict):
        strings = []
        for item in value.values():
            strings.extend(value_strings(item))
        return strings
    return [str(value)]


def operation_target(operation: dict[str, Any]) -> dict[str, str]:
    target = operation.get("target")
    if not isinstance(target, dict):
        raise ValueError("operation target must be an object with target_id/target_name")
    target_id = target.get("target_id")
    target_name = target.get("target_name")
    if not isinstance(target_id, str) or not target_id.strip():
        raise ValueError("operation target missing target_id")
    if not isinstance(target_name, str) or not target_name.strip():
        raise ValueError("operation target missing target_name")
    return {"target_id": target_id.strip(), "target_name": target_name.strip()}


def distinct_targets(sample: dict[str, Any]) -> list[dict[str, str]]:
    seen: set[str] = set()
    targets: list[dict[str, str]] = []
    operations = sample.get("operations")
    if not isinstance(operations, list):
        return targets
    for operation in operations:
        if not isinstance(operation, dict):
            continue
        target = operation_target(operation)
        if target["target_id"] in seen:
            continue
        seen.add(target["target_id"])
        targets.append(target)
    return targets


def forbidden_values_for_sample(sample: dict[str, Any]) -> set[str]:
    values: set[str] = set()
    for operation in sample.get("operations") or []:
        if not isinstance(operation, dict):
            continue
        for key in ("old_value", "new_value"):
            for value in value_strings(operation.get(key)):
                values.add(normalize_text(value))
    return {value for value in values if value}


def forgotten_values_for_sample(sample: dict[str, Any]) -> set[str]:
    values: set[str] = set()
    for operation in sample.get("operations") or []:
        if not isinstance(operation, dict):
            continue
        if str(operation.get("type", "")).lower() != "forget":
            continue
        for value in value_strings(operation.get("old_value")):
            values.add(normalize_text(value))
    return {value for value in values if value}


def final_confirmed_values_by_target(sample: dict[str, Any]) -> dict[str, str]:
    final_values: dict[str, str] = {}
    for operation in sample.get("operations") or []:
        if not isinstance(operation, dict):
            continue
        validity = str(operation.get("validity") or "confirmed").lower()
        if validity != "confirmed":
            continue
        op_type = str(operation.get("type") or "").lower()
        target_id = operation_target(operation)["target_id"]
        if op_type == "forget":
            final_values.pop(target_id, None)
            continue
        values = value_strings(operation.get("new_value"))
        if values:
            final_values[target_id] = values[0]
    return final_values


def retained_target_ids(sample: dict[str, Any]) -> set[str]:
    return set(final_confirmed_values_by_target(sample))


def operation_values_for_specs(sample: dict[str, Any]) -> list[dict[str, Any]]:
    operation_type = str(sample.get("operation_type") or "")
    retained_ids = retained_target_ids(sample)
    specs: list[dict[str, Any]] = []
    for operation in sample.get("operations") or []:
        if not isinstance(operation, dict):
            continue
        target = operation_target(operation)
        op_type = str(operation.get("type") or "").lower()
        validity = str(operation.get("validity") or "confirmed").lower()
        values = value_strings(operation.get("new_value"))
        old_values = value_strings(operation.get("old_value"))
        if operation_type == "Forget":
            if target["target_id"] in retained_ids and values:
                specs.append(
                    {
                        **target,
                        "reference_value": values[0],
                        "distractor_type": "retained_neighbor",
                        "recency_trap_candidate": False,
                        "reason": "retained neighbor target",
                    }
                )
            continue
        if operation_type == "Update":
            if values:
                dtype = "tentative_trap_neighbor" if validity in {"tentative", "retracted"} else "near_miss_third_party"
                specs.append(
                    {
                        **target,
                        "reference_value": values[0],
                        "distractor_type": dtype,
                        "recency_trap_candidate": validity in {"tentative", "retracted"},
                        "reason": f"{validity} new/current value",
                    }
                )
            for old in old_values:
                specs.append(
                    {
                        **target,
                        "reference_value": old,
                        "distractor_type": "stale_neighbor",
                        "recency_trap_candidate": True,
                        "reason": "stale old value",
                    }
                )
            continue
        if operation_type == "Reflect":
            if values:
                specs.append(
                    {
                        **target,
                        "reference_value": values[0],
                        "distractor_type": "assistant_echo_wrong_value",
                        "recency_trap_candidate": False,
                        "reason": "unconfirmed assistant abstraction",
                    }
                )
            continue
        if operation_type == "TrajectoryOps":
            for old in old_values:
                specs.append(
                    {
                        **target,
                        "reference_value": old,
                        "distractor_type": "stale_neighbor",
                        "recency_trap_candidate": True,
                        "reason": "trajectory stale value",
                    }
                )
            if values:
                dtype = "tentative_trap_neighbor" if validity in {"tentative", "retracted"} else "near_miss_third_party"
                specs.append(
                    {
                        **target,
                        "reference_value": values[0],
                        "distractor_type": dtype,
                        "recency_trap_candidate": validity in {"tentative", "retracted"},
                        "reason": f"trajectory {validity} value",
                    }
                )
            continue
        if op_type == "remember" and values:
            specs.append(
                {
                    **target,
                    "reference_value": values[0],
                    "distractor_type": "third_party_same_target",
                    "recency_trap_candidate": False,
                    "reason": "remembered target neighbor",
                }
            )
    return specs


def build_distractor_specs(sample: dict[str, Any], *, max_per_sample: int = DEFAULT_MAX_PER_SAMPLE) -> list[dict[str, Any]]:
    seen: set[tuple[str, str, str]] = set()
    specs: list[dict[str, Any]] = []
    for spec in operation_values_for_specs(sample):
        reference_value = str(spec.get("reference_value") or "").strip()
        if not reference_value:
            continue
        key = (str(spec["target_id"]), str(spec["distractor_type"]), normalize_text(reference_value))
        if key in seen:
            continue
        seen.add(key)
        specs.append(spec)
        if len(specs) >= max_per_sample:
            break
    return specs


def nearby_value(reference_value: str, forbidden_values: set[str]) -> str:
    value = reference_value.strip()
    if re.fullmatch(r"\d{4,}", value):
        base = int(value)
        for offset in (1, -1, 2, -2, 10, -10):
            candidate = str(base + offset).zfill(len(value))
            if normalize_text(candidate) not in forbidden_values:
                return candidate
    if re.fullmatch(r"\d+(?:\.\d+)?", value):
        candidate = f"{value}0"
        if normalize_text(candidate) not in forbidden_values:
            return candidate
    candidate = f"{value} (neighbor record)"
    if normalize_text(candidate) not in forbidden_values:
        return candidate
    return f"nearby {value}"


def deterministic_distractor_from_spec(
    spec: dict[str, Any],
    *,
    distractor_id: str,
    forbidden_values: set[str],
) -> dict[str, Any]:
    reference_value = str(spec.get("reference_value") or "").strip()
    distractor_type = str(spec.get("distractor_type") or "third_party_same_target")
    distractor_value = nearby_value(reference_value, forbidden_values)
    if distractor_type == "assistant_echo_wrong_value":
        dialogue = [
            {
                "role": "user",
                "content": f"I am reviewing a related note and have not confirmed any change to my own {spec['target_name']}.",
            },
            {
                "role": "assistant",
                "content": f"An assistant-only draft mentioned {distractor_value}, but I will not treat it as your confirmed memory.",
            },
        ]
        span_role = "assistant"
        span_quote = distractor_value
    elif distractor_type == "hypothetical_self":
        dialogue = [
            {
                "role": "user",
                "content": f"I almost used {distractor_value} for a similar {spec['target_name']} case, but I did not adopt it.",
            },
            {
                "role": "assistant",
                "content": "Understood; that hypothetical option is not your current stored state.",
            },
        ]
        span_role = "user"
        span_quote = distractor_value
    else:
        dialogue = [
            {
                "role": "user",
                "content": f"My coworker has a similar {spec['target_name']} value: {distractor_value}.",
            },
            {
                "role": "assistant",
                "content": "I will keep that scoped to your coworker, not you.",
            },
        ]
        span_role = "user"
        span_quote = distractor_value
    return {
        "distractor_id": distractor_id,
        "target_id": spec["target_id"],
        "target_name": spec["target_name"],
        "distractor_type": distractor_type,
        "distractor_value": distractor_value,
        "gold_values_at_generation": sorted(forbidden_values),
        "state_neutrality_rationale": "The dialogue scopes this value to a third party, assistant-only draft, or hypothetical option; it is not confirmed for the benchmark user.",
        "confusability_rationale": f"Same target/topic as {spec['target_name']} and close to reference value {reference_value}.",
        "recency_trap_candidate": bool(spec.get("recency_trap_candidate")),
        "distractor_span": {"turn_index": 1, "role": span_role, "quote": span_quote},
        "dialogue": dialogue,
    }


def local_distractor_gate_issues(
    payload: dict[str, Any],
    *,
    forbidden_values: set[str],
    forgotten_values: set[str],
    operation_type: str,
) -> list[str]:
    issues: list[str] = []
    normalized_forbidden_values = {normalize_text(value) for value in forbidden_values if normalize_text(value)}
    normalized_forgotten_values = {normalize_text(value) for value in forgotten_values if normalize_text(value)}
    distractors = payload.get("distractors")
    if not isinstance(distractors, list):
        return ["distractors must be a list."]
    for index, distractor in enumerate(distractors, start=1):
        if not isinstance(distractor, dict):
            issues.append(f"distractor #{index} is not an object.")
            continue
        did = str(distractor.get("distractor_id") or f"#{index}")
        dtype = str(distractor.get("distractor_type") or "")
        if dtype not in DISTRACTOR_TYPES:
            issues.append(f"{did} has unsupported distractor_type: {dtype!r}.")
        dialogue = distractor.get("dialogue")
        if not isinstance(dialogue, list) or not (2 <= len(dialogue) <= 6):
            issues.append(f"{did} dialogue must contain 2-6 messages.")
            continue
        if not isinstance(dialogue[0], dict) or dialogue[0].get("role") != "user":
            issues.append(f"{did} dialogue must start with user.")
        if not isinstance(dialogue[-1], dict) or dialogue[-1].get("role") != "assistant":
            issues.append(f"{did} dialogue must end with assistant.")
        for previous, current in zip(dialogue, dialogue[1:]):
            if not isinstance(previous, dict) or not isinstance(current, dict):
                issues.append(f"{did} dialogue messages must be objects.")
                continue
            if previous.get("role") == current.get("role"):
                issues.append(f"{did} dialogue roles must alternate.")
        value = str(distractor.get("distractor_value") or "").strip()
        normalized_value = normalize_text(value)
        if not value:
            issues.append(f"{did} missing distractor_value.")
        # Structural-only: require the span object to exist; whether the quote is
        # verbatim and whether it actually supports the value are SEMANTIC checks
        # delegated to the LLM verifier (span_valid). Exact-substring matching here
        # produced false negatives on capitalization / pronoun / phrasing differences.
        span = distractor.get("distractor_span")
        span_role = ""
        if isinstance(span, dict):
            span_role = str(span.get("role") or "")
        else:
            issues.append(f"{did} missing distractor_span object.")
        # Leak guards: a distractor must never expose the benchmark user's real
        # forgotten or forbidden gold value as their own.
        if normalized_value in normalized_forgotten_values:
            issues.append(f"{did} leaks a forgotten raw value.")
        if normalized_value in normalized_forbidden_values:
            if dtype not in ASSISTANT_ONLY_TYPES or span_role != "assistant":
                issues.append(f"{did} uses a forbidden value without assistant-only containment.")
    return issues


def build_generation_prompt(
    *,
    template: str,
    sample: dict[str, Any],
    source_file: str,
    specs: list[dict[str, Any]],
    forbidden_values: set[str],
    forgotten_values: set[str],
    retry_feedback: str = "",
) -> str:
    return (
        template.replace("{{SOURCE_FILE}}", source_file)
        .replace("{{OPERATION_TYPE}}", str(sample.get("operation_type") or ""))
        .replace("{{SAMPLE_JSON}}", json.dumps(sample, ensure_ascii=False, indent=2))
        .replace("{{DISTRACTOR_SPECS}}", json.dumps(specs, ensure_ascii=False, indent=2))
        .replace("{{FORBIDDEN_VALUES}}", json.dumps(sorted(forbidden_values), ensure_ascii=False))
        .replace("{{FORGOTTEN_VALUES}}", json.dumps(sorted(forgotten_values), ensure_ascii=False))
        .replace("{{RETRY_FEEDBACK}}", retry_feedback or "None.")
    )


def build_verify_prompt(*, template: str, payload: dict[str, Any], source_file: str) -> str:
    return (
        template.replace("{{SOURCE_FILE}}", source_file)
        .replace("{{DISTRACTOR_JSON}}", json.dumps(payload, ensure_ascii=False, indent=2))
    )


def deterministic_payload_for_sample(
    sample: dict[str, Any],
    *,
    source_file: str,
    max_per_sample: int,
) -> dict[str, Any]:
    forbidden = forbidden_values_for_sample(sample)
    specs = build_distractor_specs(sample, max_per_sample=max_per_sample)
    distractors = [
        deterministic_distractor_from_spec(
            spec,
            distractor_id=f"d{index}",
            forbidden_values=forbidden,
        )
        for index, spec in enumerate(specs, start=1)
    ]
    for distractor in distractors:
        distractor["attempts_used"] = 0
        distractor["verify_verdict"] = "LOCAL"
    return {
        "source_file": source_file,
        "operation_type": sample.get("operation_type"),
        "distractors": distractors,
        "skipped_targets": [],
    }


def normalize_generated_payload(payload: dict[str, Any], *, source_file: str, operation_type: str, attempt: int) -> dict[str, Any]:
    distractors = payload.get("distractors")
    if not isinstance(distractors, list):
        distractors = []
    normalized = {
        "source_file": source_file,
        "operation_type": operation_type,
        "distractors": distractors,
        "skipped_targets": payload.get("skipped_targets", []),
    }
    for index, distractor in enumerate(normalized["distractors"], start=1):
        if not isinstance(distractor, dict):
            continue
        distractor.setdefault("distractor_id", f"d{index}")
        distractor["attempts_used"] = attempt
    return normalized


def verify_payload_with_llm(
    payload: dict[str, Any],
    *,
    verifier_caller: LlmCaller | None,
    verify_model: str,
    verify_template: str,
    source_file: str,
) -> tuple[bool, str, dict[str, Any] | None]:
    if verifier_caller is None:
        return True, "No verifier configured.", None
    try:
        verifier_text = verifier_caller(
            verify_model,
            build_verify_prompt(template=verify_template, payload=payload, source_file=source_file),
        )
    except Exception as exc:  # noqa: BLE001 - a verifier transport/service error
        # must not crash the whole batch. Treat it as a (recorded) attempt failure
        # so other samples still complete and the cause is visible in failed_attempts.
        return False, f"verifier call failed (infrastructure, not a content issue): {exc}", None
    parsed = parse_json_object(verifier_text)
    if not parsed:
        return False, f"Could not parse verifier JSON: {verifier_text[:300]}", None
    verdict = str(parsed.get("verdict") or "").upper()
    passed = verdict == "PASS" and all(
        bool(parsed.get(key))
        for key in ("state_neutral", "confusable", "type_matches", "span_valid")
    )
    return passed, str(parsed.get("reason") or verifier_text[:300]), parsed


def generate_distractors_for_sample(
    sample: dict[str, Any],
    *,
    source_file: str,
    llm_caller: LlmCaller | None = call_llm if USE_LLM_BY_DEFAULT else None,
    verifier_caller: LlmCaller | None = call_llm if USE_LLM_BY_DEFAULT else None,
    generation_model: str = DEFAULT_GENERATION_MODEL,
    verify_model: str = DEFAULT_VERIFY_MODEL,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
    max_per_sample: int = DEFAULT_MAX_PER_SAMPLE,
    generation_template: str | None = None,
    verify_template: str | None = None,
) -> dict[str, Any]:
    operation_type = str(sample.get("operation_type") or "")
    forbidden = forbidden_values_for_sample(sample)
    forgotten = forgotten_values_for_sample(sample)
    specs = build_distractor_specs(sample, max_per_sample=max_per_sample)
    if llm_caller is None:
        payload = deterministic_payload_for_sample(
            sample,
            source_file=source_file,
            max_per_sample=max_per_sample,
        )
        issues = local_distractor_gate_issues(
            payload,
            forbidden_values=forbidden,
            forgotten_values=forgotten,
            operation_type=operation_type,
        )
        if issues:
            payload["skipped_targets"] = [
                {"reason": "deterministic fallback failed local gate", "issues": issues}
            ]
            payload["distractors"] = []
        return payload

    generation_template = generation_template or read_text(GENERATION_TEMPLATE_PATH)
    verify_template = verify_template or read_text(VERIFY_TEMPLATE_PATH)
    retry_feedback = ""
    last_issues: list[str] = []
    failed_attempts: list[dict[str, Any]] = []
    for attempt in range(1, max_attempts + 1):
        effective_feedback = (
            f"{retry_feedback}\n\n{RETRY_CONSTRAINTS_REMINDER}" if retry_feedback else ""
        )
        prompt = build_generation_prompt(
            template=generation_template,
            sample=sample,
            source_file=source_file,
            specs=specs,
            forbidden_values=forbidden,
            forgotten_values=forgotten,
            retry_feedback=effective_feedback,
        )
        raw_text = llm_caller(generation_model, prompt)
        parsed = parse_json_object(raw_text)
        if not parsed:
            issues = [f"Attempt {attempt}: output was not valid JSON."]
            retry_feedback = issues[0]
            last_issues = issues
            failed_attempts.append(
                {
                    "attempt": attempt,
                    "stage": "json_parse",
                    "issues": issues,
                    "payload": None,
                    "raw_text": raw_text,
                }
            )
            continue
        payload = normalize_generated_payload(
            parsed,
            source_file=source_file,
            operation_type=operation_type,
            attempt=attempt,
        )
        issues = local_distractor_gate_issues(
            payload,
            forbidden_values=forbidden,
            forgotten_values=forgotten,
            operation_type=operation_type,
        )
        if issues:
            retry_feedback = "Attempt {} local gate failed:\n{}".format(
                attempt,
                "\n".join(f"- {issue}" for issue in issues),
            )
            last_issues = issues
            failed_attempts.append(
                {
                    "attempt": attempt,
                    "stage": "local_gate",
                    "issues": issues,
                    "payload": payload,
                }
            )
            continue
        verified, verify_reason, verify_payload = verify_payload_with_llm(
            payload,
            verifier_caller=verifier_caller,
            verify_model=verify_model,
            verify_template=verify_template,
            source_file=source_file,
        )
        if not verified:
            issues = [f"Attempt {attempt} verifier failed: {verify_reason}"]
            retry_feedback = issues[0]
            last_issues = issues
            failed_attempts.append(
                {
                    "attempt": attempt,
                    "stage": "verifier",
                    "issues": issues,
                    "payload": payload,
                    "verify_reason": verify_reason,
                    "verify_payload": verify_payload,
                }
            )
            continue
        for distractor in payload["distractors"]:
            if isinstance(distractor, dict):
                distractor["verify_verdict"] = (
                    str((verify_payload or {}).get("verdict") or "PASS").upper()
                    if verify_payload
                    else "PASS"
                )
                distractor["verify_reason"] = verify_reason
        if failed_attempts:
            payload["failed_attempts"] = failed_attempts
        return payload
    return {
        "source_file": source_file,
        "operation_type": operation_type,
        "distractors": [],
        "failed_attempts": failed_attempts,
        "skipped_targets": [
            {
                "reason": "max attempts exhausted",
                "issues": last_issues,
                "spec_count": len(specs),
                "attempts": failed_attempts,
            }
        ],
    }


def process_one_file(
    evidence_path: Path,
    *,
    output_dir: Path,
    llm_caller: LlmCaller | None,
    verifier_caller: LlmCaller | None,
    generation_model: str,
    verify_model: str,
    max_attempts: int,
    max_per_sample: int,
    generation_template: str | None,
    verify_template: str | None,
) -> Path:
    sample = read_json(evidence_path)
    if not isinstance(sample, dict):
        raise ValueError(f"Evidence file must contain a JSON object: {evidence_path}")
    payload = generate_distractors_for_sample(
        sample,
        source_file=evidence_path.name,
        llm_caller=llm_caller,
        verifier_caller=verifier_caller,
        generation_model=generation_model,
        verify_model=verify_model,
        max_attempts=max_attempts,
        max_per_sample=max_per_sample,
        generation_template=generation_template,
        verify_template=verify_template,
    )
    output_path = output_dir / evidence_path.name
    write_json(output_path, payload)
    # Persist each failed generation attempt for post-hoc debugging. These
    # *_attempt_N.json files live in the distractor dir but are ignored by the
    # injection step, which only loads the exact evidence file name.
    stem = evidence_path.stem
    for record in payload.get("failed_attempts", []):
        if not isinstance(record, dict):
            continue
        attempt_index = record.get("attempt")
        if attempt_index is None:
            continue
        attempt_payload = {
            "source_file": evidence_path.name,
            "operation_type": payload.get("operation_type"),
            **record,
        }
        write_json(output_dir / f"{stem}_attempt_{attempt_index}.json", attempt_payload)
    return output_path


def process_evidence_directory(
    *,
    evidence_dir: Path,
    output_dir: Path,
    llm_caller: LlmCaller | None = call_llm if USE_LLM_BY_DEFAULT else None,
    verifier_caller: LlmCaller | None = call_llm if USE_LLM_BY_DEFAULT else None,
    generation_model: str = DEFAULT_GENERATION_MODEL,
    verify_model: str = DEFAULT_VERIFY_MODEL,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
    max_per_sample: int = DEFAULT_MAX_PER_SAMPLE,
    concurrency: int = DEFAULT_CONCURRENCY,
    show_progress: bool = True,
    progress_factory: ProgressFactory = tqdm,
) -> list[Path]:
    evidence_files = iter_evidence_files(evidence_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    pending_files = [
        path for path in evidence_files if not (output_dir / path.name).exists()
    ]
    existing_outputs = [
        output_dir / path.name for path in evidence_files if (output_dir / path.name).exists()
    ]
    generation_template = read_text(GENERATION_TEMPLATE_PATH) if llm_caller else None
    verify_template = read_text(VERIFY_TEMPLATE_PATH) if verifier_caller else None
    if not pending_files:
        return existing_outputs
    if concurrency <= 1 or len(pending_files) <= 1:
        outputs: list[Path] = existing_outputs[:]
        for path in progress_factory(
            pending_files,
            desc="Generating distractors",
            unit="file",
            disable=not show_progress,
        ):
            outputs.append(
                process_one_file(
                    path,
                    output_dir=output_dir,
                    llm_caller=llm_caller,
                    verifier_caller=verifier_caller,
                    generation_model=generation_model,
                    verify_model=verify_model,
                    max_attempts=max_attempts,
                    max_per_sample=max_per_sample,
                    generation_template=generation_template,
                    verify_template=verify_template,
                )
            )
        return outputs

    outputs_by_path: dict[Path, Path] = {}
    with ThreadPoolExecutor(max_workers=concurrency) as executor:
        future_to_path = {
            executor.submit(
                process_one_file,
                path,
                output_dir=output_dir,
                llm_caller=llm_caller,
                verifier_caller=verifier_caller,
                generation_model=generation_model,
                verify_model=verify_model,
                max_attempts=max_attempts,
                max_per_sample=max_per_sample,
                generation_template=generation_template,
                verify_template=verify_template,
            ): path
            for path in pending_files
        }
        for future in progress_factory(
            as_completed(future_to_path),
            desc="Generating distractors",
            unit="file",
            total=len(future_to_path),
            disable=not show_progress,
        ):
            path = future_to_path[future]
            outputs_by_path[path] = future.result()
    generated_outputs = [outputs_by_path[path] for path in pending_files]
    return existing_outputs + generated_outputs


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate v24 distractor conversations.")
    parser.add_argument("--evidence-dir", type=Path, default=Path(os.getenv("DISTRACTOR_EVIDENCE_DIR", str(DEFAULT_EVIDENCE_DIR))))
    parser.add_argument("--output-dir", type=Path, default=Path(os.getenv("DISTRACTOR_OUTPUT_DIR", str(DEFAULT_OUTPUT_DIR))))
    parser.add_argument("--generation-model", default=DEFAULT_GENERATION_MODEL)
    parser.add_argument("--verify-model", default=DEFAULT_VERIFY_MODEL)
    parser.add_argument("--max-attempts", type=int, default=DEFAULT_MAX_ATTEMPTS)
    parser.add_argument("--max-per-sample", type=int, default=DEFAULT_MAX_PER_SAMPLE)
    parser.add_argument("--concurrency", type=int, default=DEFAULT_CONCURRENCY)
    parser.add_argument(
        "--disable-llm",
        action="store_true",
        help="Use deterministic local distractor generation without API calls.",
    )
    return parser.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> int:
    args = parse_args(argv)
    use_llm = USE_LLM_BY_DEFAULT and not args.disable_llm
    outputs = process_evidence_directory(
        evidence_dir=args.evidence_dir,
        output_dir=args.output_dir,
        llm_caller=call_llm if use_llm else None,
        verifier_caller=call_llm if use_llm else None,
        generation_model=args.generation_model,
        verify_model=args.verify_model,
        max_attempts=args.max_attempts,
        max_per_sample=args.max_per_sample,
        concurrency=args.concurrency,
    )
    print(f"Saved {len(outputs)} distractor file(s) to {args.output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())