"""Leak-free controller transition dataset for Gate 5A."""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from .bytecode import CognitiveOp
from .consolidation import ConsolidatedMemoryRecord, MemoryCard, agent_safe_memory_cards
from .ledger import TraceLedger, TraceLedgerRecord, canonical_json, stable_hash
from .trace import utc_now
from .workbench import oracle_leak_keys

CONTROLLER_DATASET_SCHEMA_VERSION = 1
DATASET_FILE = "controller_transitions.jsonl"
SPLIT_FILES = {
    "train": "train.jsonl",
    "validation": "validation.jsonl",
    "test": "test.jsonl",
    "adversarial": "adversarial.jsonl",
    "negative": "negative.jsonl",
}


@dataclass(slots=True)
class ControllerTransition:
    metadata: dict[str, Any]
    model_input: dict[str, Any]
    labels: dict[str, Any]

    @property
    def transition_id(self) -> str:
        return str(self.metadata["transition_id"])

    def stable_payload(self) -> dict[str, Any]:
        return {
            "metadata": {key: value for key, value in self.metadata.items() if key != "transition_id"},
            "model_input": self.model_input,
            "labels": self.labels,
        }

    def stable_hash(self) -> str:
        return stable_hash(self.stable_payload())

    def to_dict(self) -> dict[str, Any]:
        return {
            "metadata": self.metadata,
            "model_input": self.model_input,
            "labels": self.labels,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ControllerTransition":
        return cls(
            metadata=dict(data["metadata"]),
            model_input=dict(data["model_input"]),
            labels=dict(data["labels"]),
        )


@dataclass(slots=True)
class ControllerDatasetManifest:
    schema_version: int
    output_dir: str
    dataset_path: str
    split_paths: dict[str, str]
    transition_count: int
    split_counts: dict[str, int]
    action_distribution: dict[str, int]
    distinct_target_actions: int
    leak_audit: dict[str, Any]
    canary_audit: dict[str, Any]
    replay_validation: dict[str, Any]
    forensics: dict[str, Any]
    sequence_shortcut_audit: dict[str, Any]
    grounding_context_audit: dict[str, Any]
    offline_baselines: dict[str, Any]
    source_trace_count: int
    source_trace_hashes: list[str]
    source_gate_counts: dict[str, int]
    build_timestamp: str = field(default_factory=utc_now)

    def to_dict(self) -> dict[str, Any]:
        return {
            "schema_version": self.schema_version,
            "output_dir": self.output_dir,
            "dataset_path": self.dataset_path,
            "split_paths": self.split_paths,
            "transition_count": self.transition_count,
            "split_counts": self.split_counts,
            "action_distribution": self.action_distribution,
            "distinct_target_actions": self.distinct_target_actions,
            "leak_audit": self.leak_audit,
            "canary_audit": self.canary_audit,
            "replay_validation": self.replay_validation,
            "forensics": self.forensics,
            "sequence_shortcut_audit": self.sequence_shortcut_audit,
            "grounding_context_audit": self.grounding_context_audit,
            "offline_baselines": self.offline_baselines,
            "source_trace_count": self.source_trace_count,
            "source_trace_hashes": self.source_trace_hashes,
            "source_gate_counts": self.source_gate_counts,
            "build_timestamp": self.build_timestamp,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ControllerDatasetManifest":
        return cls(
            schema_version=int(data["schema_version"]),
            output_dir=str(data["output_dir"]),
            dataset_path=str(data["dataset_path"]),
            split_paths={str(key): str(value) for key, value in dict(data["split_paths"]).items()},
            transition_count=int(data["transition_count"]),
            split_counts={str(key): int(value) for key, value in dict(data["split_counts"]).items()},
            action_distribution={str(key): int(value) for key, value in dict(data["action_distribution"]).items()},
            distinct_target_actions=int(data["distinct_target_actions"]),
            leak_audit=dict(data["leak_audit"]),
            canary_audit=dict(data.get("canary_audit") or {}),
            replay_validation=dict(data.get("replay_validation") or {}),
            forensics=dict(data.get("forensics") or {}),
            sequence_shortcut_audit=dict(data.get("sequence_shortcut_audit") or {}),
            grounding_context_audit=dict(data.get("grounding_context_audit") or {}),
            offline_baselines=dict(data["offline_baselines"]),
            source_trace_count=int(data["source_trace_count"]),
            source_trace_hashes=[str(item) for item in data.get("source_trace_hashes", [])],
            source_gate_counts={str(key): int(value) for key, value in dict(data.get("source_gate_counts") or {}).items()},
            build_timestamp=str(data.get("build_timestamp") or utc_now()),
        )


@dataclass(slots=True)
class DatasetValidationReport:
    invalid_transition_count: int = 0
    oracle_leak_count: int = 0
    unsupported_source_count: int = 0
    action_mask_error_count: int = 0
    canary_terms_tested: int = 0
    canary_leak_count: int = 0
    hidden_value_leak_count: int = 0
    oracle_field_leak_count: int = 0
    invalid_arg_count: int = 0
    illegal_action_count: int = 0
    action_mask_violation_count: int = 0
    unsafe_action_count: int = 0
    target_action_count: int = 0
    target_action_legal_count: int = 0
    invalid_transition_reports: list[dict[str, Any]] = field(default_factory=list)

    @property
    def target_action_legal_rate(self) -> float:
        if self.target_action_count == 0:
            return 0.0
        return self.target_action_legal_count / self.target_action_count

    def to_dict(self) -> dict[str, Any]:
        return {
            "invalid_transition_count": self.invalid_transition_count,
            "oracle_leak_count": self.oracle_leak_count,
            "unsupported_source_count": self.unsupported_source_count,
            "action_mask_error_count": self.action_mask_error_count,
            "canary_terms_tested": self.canary_terms_tested,
            "canary_leak_count": self.canary_leak_count,
            "hidden_value_leak_count": self.hidden_value_leak_count,
            "oracle_field_leak_count": self.oracle_field_leak_count,
            "invalid_arg_count": self.invalid_arg_count,
            "illegal_action_count": self.illegal_action_count,
            "action_mask_violation_count": self.action_mask_violation_count,
            "unsafe_action_count": self.unsafe_action_count,
            "target_action_legal_rate": self.target_action_legal_rate,
            "invalid_transition_reports": self.invalid_transition_reports,
        }


class ControllerDatasetValidator:
    def __init__(self, source_traces: list[TraceLedgerRecord]) -> None:
        self.source_hashes = {trace.trace_id: trace.content_hash for trace in source_traces if trace.verify_hash()}
        self.legal_actions = {op.value for op in CognitiveOp}

    def validate(self, transitions: list[ControllerTransition]) -> DatasetValidationReport:
        report = DatasetValidationReport(target_action_count=len(transitions))
        for transition in transitions:
            reasons: list[str] = []
            leaks = oracle_leak_keys(transition.model_input)
            oracle_leaked = False
            if leaks:
                oracle_leaked = True
                report.oracle_field_leak_count += 1
                reasons.append("oracle_leak:" + ",".join(leaks))
            trace_id = str(transition.metadata.get("trace_id", ""))
            trace_hash = str(transition.metadata.get("trace_hash", ""))
            if self.source_hashes.get(trace_id) != trace_hash:
                report.unsupported_source_count += 1
                reasons.append("unsupported_source")
            mask = dict(transition.model_input.get("action_mask") or {})
            legal_actions = set(transition.model_input.get("legal_actions") or [])
            target_action = str(transition.labels.get("target_action", ""))
            target_args = transition.labels.get("target_action_args")
            if not isinstance(target_args, dict) or not _json_safe(target_args):
                report.invalid_arg_count += 1
                reasons.append("invalid_args")
            if legal_actions != self.legal_actions or set(mask) != self.legal_actions:
                report.action_mask_error_count += 1
                report.action_mask_violation_count += 1
                reasons.append("mask_shape")
            elif not all(isinstance(value, bool) for value in mask.values()):
                report.action_mask_error_count += 1
                report.action_mask_violation_count += 1
                reasons.append("mask_values")
            elif target_action not in self.legal_actions:
                report.action_mask_error_count += 1
                report.action_mask_violation_count += 1
                report.illegal_action_count += 1
                report.unsafe_action_count += 1
                reasons.append("illegal_target")
            elif not bool(mask.get(target_action)):
                report.action_mask_error_count += 1
                report.action_mask_violation_count += 1
                reasons.append("target_not_enabled")
            else:
                report.target_action_legal_count += 1
            forbidden_values = _forbidden_metadata_values(dict(transition.metadata.get("evaluator_metadata") or {}))
            report.canary_terms_tested += len(set(forbidden_values))
            forbidden_hits = _forbidden_value_hits(transition.model_input, forbidden_values)
            if forbidden_hits:
                oracle_leaked = True
                report.canary_leak_count += 1
                report.hidden_value_leak_count += 1
                reasons.append("forbidden_evaluator_value:" + ",".join(forbidden_hits))
            if oracle_leaked:
                report.oracle_leak_count += 1
            if reasons:
                report.invalid_transition_count += 1
                report.invalid_transition_reports.append(
                    {"transition_id": transition.transition_id, "reasons": reasons}
                )
        return report


def build_controller_dataset(state_dir: Path, output_dir: Path) -> ControllerDatasetManifest:
    state_dir = Path(state_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    ledger_path = _ensure_gate4_ledger(state_dir)
    ledger = TraceLedger(ledger_path.parent, file_name=ledger_path.name)
    traces = ledger.read_all()
    memories = _load_consolidated_memories(ledger_path.parent)
    safe_cards = agent_safe_memory_cards(memories)
    transitions = _build_transitions(traces, safe_cards)
    validation = ControllerDatasetValidator(traces).validate(transitions)
    offline_baselines = evaluate_offline_baselines(transitions)
    sequence_shortcut_audit = evaluate_sequence_shortcuts(transitions)

    dataset_path = output_dir / DATASET_FILE
    _write_jsonl(dataset_path, [transition.to_dict() for transition in transitions])
    split_paths: dict[str, str] = {}
    for split, file_name in SPLIT_FILES.items():
        split_path = output_dir / file_name
        split_paths[split] = str(split_path)
        _write_jsonl(split_path, [transition.to_dict() for transition in transitions if _split(transition) == split])
    invalid_path = output_dir / "invalid_transitions.json"
    invalid_path.write_text(json.dumps(validation.invalid_transition_reports, indent=2, sort_keys=True), encoding="utf-8")

    split_counts = {split: sum(1 for transition in transitions if _split(transition) == split) for split in SPLIT_FILES}
    action_distribution: dict[str, int] = {}
    for transition in transitions:
        action = str(transition.labels["target_action"])
        action_distribution[action] = action_distribution.get(action, 0) + 1
    manifest = ControllerDatasetManifest(
        schema_version=CONTROLLER_DATASET_SCHEMA_VERSION,
        output_dir=str(output_dir),
        dataset_path=str(dataset_path),
        split_paths=split_paths,
        transition_count=len(transitions),
        split_counts=split_counts,
        action_distribution=dict(sorted(action_distribution.items())),
        distinct_target_actions=len(action_distribution),
        leak_audit=validation.to_dict(),
        canary_audit=_canary_audit(validation),
        replay_validation={
            "replay_validation_passed": False,
            "not_run": True,
            "transition_id_match": False,
            "transition_hash_match": False,
            "source_hash_match": False,
            "rebuild_transition_count": len(transitions),
            "original_transition_count": len(transitions),
        },
        forensics=_forensics_summary(
            transitions=transitions,
            validation=validation,
            source_trace_count=len(traces),
            split_counts=split_counts,
            distinct_action_count=len(action_distribution),
            offline_baselines=offline_baselines,
        ),
        sequence_shortcut_audit=sequence_shortcut_audit,
        grounding_context_audit=_grounding_context_audit(transitions, sequence_shortcut_audit),
        offline_baselines=offline_baselines,
        source_trace_count=len(traces),
        source_trace_hashes=[trace.content_hash for trace in traces],
        source_gate_counts=_source_gate_counts(transitions),
    )
    (output_dir / "manifest.json").write_text(json.dumps(manifest.to_dict(), indent=2, sort_keys=True), encoding="utf-8")
    return manifest


def replay_validate_controller_dataset(
    state_dir: Path,
    manifest: ControllerDatasetManifest,
    replay_output_dir: Path,
) -> dict[str, Any]:
    original = _read_transitions(Path(manifest.dataset_path))
    rebuilt = build_controller_dataset(Path(state_dir), Path(replay_output_dir))
    replayed = _read_transitions(Path(rebuilt.dataset_path))
    original_ids = [transition.transition_id for transition in original]
    replayed_ids = [transition.transition_id for transition in replayed]
    original_hashes = [transition.stable_hash() for transition in original]
    replayed_hashes = [transition.stable_hash() for transition in replayed]
    source_hash_match = list(manifest.source_trace_hashes) == list(rebuilt.source_trace_hashes)
    transition_id_match = original_ids == replayed_ids
    transition_hash_match = original_hashes == replayed_hashes
    return {
        "replay_validation_passed": transition_id_match and transition_hash_match and source_hash_match,
        "transition_id_match": transition_id_match,
        "transition_hash_match": transition_hash_match,
        "source_hash_match": source_hash_match,
        "rebuild_transition_count": len(replayed),
        "original_transition_count": len(original),
        "replay_dataset_path": rebuilt.dataset_path,
    }


def evaluate_offline_baselines(transitions: list[ControllerTransition]) -> dict[str, Any]:
    train = [transition for transition in transitions if _split(transition) == "train"]
    evaluation_splits = {
        split: [transition for transition in transitions if _split(transition) == split]
        for split in ["train", "validation", "test", "adversarial", "negative"]
    }
    baselines = {
        "most_frequent_action": MostFrequentActionPolicy(train),
        "legal_stage_replay": LegalStageReplayPolicy(train),
        "oracle_free_trace_bigram": OracleFreeTraceBigramPolicy(train),
    }
    results = {name: _evaluate_policy(policy, evaluation_splits) for name, policy in baselines.items()}
    validation_scores = {
        name: metrics["splits"]["validation"]["top1_accuracy"]
        for name, metrics in results.items()
    }
    best_name = sorted(validation_scores, key=lambda item: (-validation_scores[item], item))[0]
    results["best_baseline"] = {
        "name": best_name,
        "validation_top1_accuracy": validation_scores[best_name],
        "beats_most_frequent_on_validation": validation_scores[best_name] > validation_scores["most_frequent_action"],
    }
    return results


def evaluate_sequence_shortcuts(transitions: list[ControllerTransition]) -> dict[str, Any]:
    train = [transition for transition in transitions if _split(transition) == "train"]
    validation = [transition for transition in transitions if _split(transition) == "validation"]
    previous_action_accuracy = _shortcut_accuracy(
        train,
        validation,
        lambda transition: str(transition.model_input.get("previous_action") or "<START>"),
    )
    step_index_accuracy = _shortcut_accuracy(
        train,
        validation,
        lambda transition: str(transition.metadata.get("step_index", 0)),
    )
    previous_action_plus_step_accuracy = _shortcut_accuracy(
        train,
        validation,
        lambda transition: f"{transition.model_input.get('previous_action') or '<START>'}|{transition.metadata.get('step_index', 0)}",
    )
    bigram_accuracy = previous_action_accuracy
    max_accuracy = max(previous_action_accuracy, step_index_accuracy, previous_action_plus_step_accuracy, bigram_accuracy)
    shortcut_risk = "high" if max_accuracy >= 0.85 else "medium" if max_accuracy >= 0.65 else "low"
    previous_action_ambiguity = _ambiguous_key_stats(
        transitions,
        lambda transition: str(transition.model_input.get("previous_action") or "<START>"),
    )
    step_index_ambiguity = _ambiguous_key_stats(
        transitions,
        lambda transition: str(transition.metadata.get("step_index", 0)),
    )
    previous_action_plus_step_ambiguity = _ambiguous_key_stats(
        transitions,
        lambda transition: f"{transition.model_input.get('previous_action') or '<START>'}|{transition.metadata.get('step_index', 0)}",
    )
    return {
        "previous_action_accuracy": previous_action_accuracy,
        "step_index_accuracy": step_index_accuracy,
        "previous_action_plus_step_accuracy": previous_action_plus_step_accuracy,
        "trace_template_id_available": False,
        "trace_template_id_accuracy": None,
        "bigram_accuracy": bigram_accuracy,
        "ambiguous_previous_action_count": previous_action_ambiguity["ambiguous_key_count"],
        "ambiguous_step_index_count": step_index_ambiguity["ambiguous_key_count"],
        "ambiguous_previous_action_plus_step_count": previous_action_plus_step_ambiguity["ambiguous_key_count"],
        "grounding_required_transition_count": previous_action_ambiguity["ambiguous_transition_count"],
        "shortcut_risk": shortcut_risk,
    }


class MostFrequentActionPolicy:
    def __init__(self, train: list[ControllerTransition]) -> None:
        counts = _action_counts(train)
        self.action = sorted(counts, key=lambda item: (-counts[item], item))[0] if counts else CognitiveOp.ANSWER.value

    def predict(self, transition: ControllerTransition) -> str:
        return self.action


class LegalStageReplayPolicy:
    def __init__(self, train: list[ControllerTransition]) -> None:
        self.by_step: dict[int, str] = {}
        grouped: dict[int, dict[str, int]] = {}
        for transition in train:
            step = int(transition.model_input.get("recent_step_index", 0))
            action = str(transition.labels["target_action"])
            grouped.setdefault(step, {})[action] = grouped.setdefault(step, {}).get(action, 0) + 1
        for step, counts in grouped.items():
            self.by_step[step] = sorted(counts, key=lambda item: (-counts[item], item))[0]
        self.fallback = MostFrequentActionPolicy(train).action

    def predict(self, transition: ControllerTransition) -> str:
        return self.by_step.get(int(transition.model_input.get("recent_step_index", 0)), self.fallback)


class OracleFreeTraceBigramPolicy:
    def __init__(self, train: list[ControllerTransition]) -> None:
        self.by_previous: dict[str, str] = {}
        grouped: dict[str, dict[str, int]] = {}
        for transition in train:
            previous = str(transition.model_input.get("previous_action") or "<START>")
            action = str(transition.labels["target_action"])
            grouped.setdefault(previous, {})[action] = grouped.setdefault(previous, {}).get(action, 0) + 1
        for previous, counts in grouped.items():
            self.by_previous[previous] = sorted(counts, key=lambda item: (-counts[item], item))[0]
        self.fallback = MostFrequentActionPolicy(train).action

    def predict(self, transition: ControllerTransition) -> str:
        previous = str(transition.model_input.get("previous_action") or "<START>")
        return self.by_previous.get(previous, self.fallback)


def _ensure_gate4_ledger(state_dir: Path) -> Path:
    candidates = sorted(state_dir.rglob("trace_ledger.jsonl")) if state_dir.exists() else []
    if candidates:
        return candidates[-1]
    from .evaluation import run_gate4_dream_eval

    run_gate4_dream_eval(state_dir=state_dir / "gate4_seed")
    candidates = sorted(state_dir.rglob("trace_ledger.jsonl"))
    if not candidates:
        raise FileNotFoundError("Gate 5A could not seed or locate a Gate 4 trace ledger")
    return candidates[-1]


def _load_consolidated_memories(run_dir: Path) -> list[ConsolidatedMemoryRecord]:
    path = run_dir / "consolidated_memory.json"
    if not path.exists():
        return []
    return [ConsolidatedMemoryRecord.from_dict(item) for item in json.loads(path.read_text(encoding="utf-8"))]


def _build_transitions(traces: list[TraceLedgerRecord], safe_cards: list[MemoryCard]) -> list[ControllerTransition]:
    transitions: list[ControllerTransition] = []
    for trace_index, trace in enumerate(sorted(traces, key=lambda item: item.trace_id)):
        assigned_split = _assign_split(trace, trace_index)
        for step_index, action in enumerate(trace.bytecode_ops):
            previous_action = trace.bytecode_ops[step_index - 1] if step_index > 0 else None
            transition = ControllerTransition(
                metadata={
                    "schema_version": CONTROLLER_DATASET_SCHEMA_VERSION,
                    "trace_id": trace.trace_id,
                    "trace_hash": trace.content_hash,
                    "step_index": step_index,
                    "transition_id": "",
                    "gate_source": _gate_source(trace),
                    "split": assigned_split,
                    "evaluator_metadata": trace.evaluator_metadata,
                },
                model_input=_model_input(trace, safe_cards, step_index, previous_action),
                labels=_labels(trace, step_index, action),
            )
            transition.metadata["transition_id"] = "ctrl_" + transition.stable_hash()[:24]
            transitions.append(transition)
    return sorted(transitions, key=lambda item: item.transition_id)


def _model_input(
    trace: TraceLedgerRecord,
    safe_cards: list[MemoryCard],
    step_index: int,
    previous_action: str | None,
) -> dict[str, Any]:
    legal_actions = [op.value for op in CognitiveOp]
    recent_summaries = [str(item) for item in trace.agent_visible.get("step_summaries", [])[:step_index]]
    visible_tools = list(trace.agent_visible.get("visible_tool_summaries", []))
    previous_tool_summary = _previous_tool_summary(visible_tools, step_index)
    return {
        "agent_visible_observation": {
            "task_description": str(trace.agent_visible.get("task_description", "")),
            "recent_step_summaries": recent_summaries[-4:],
            "visible_tool_summaries": visible_tools[: max(1, min(len(visible_tools), step_index + 1))],
        },
        "workspace_summary": _workspace_summary(recent_summaries, step_index),
        "retrieved_memory_cards": [_agent_safe_card_dict(card, trace.evaluator_metadata) for card in safe_cards[:3]],
        "available_skill_cards": _available_skill_cards(trace),
        "previous_action": previous_action,
        "previous_tool_result_summary": previous_tool_summary,
        "legal_actions": legal_actions,
        "action_mask": {action: True for action in legal_actions},
        "verifier_context_summary": _verifier_context_summary(trace),
        "recent_step_index": step_index,
    }


def _labels(trace: TraceLedgerRecord, step_index: int, action: str) -> dict[str, Any]:
    return {
        "target_action": action,
        "target_action_args": {},
        "target_skill_id": trace.reused_skill_id if action == CognitiveOp.RETRIEVE_SKILL.value else None,
        "future_success_label": trace.success,
        "remaining_steps_label": max(0, len(trace.bytecode_ops) - step_index - 1),
        "used_fallback_later": trace.fallback_used and step_index < len(trace.bytecode_ops) - 1,
        "skill_reuse_later": bool(trace.reused_skill_id) and step_index < len(trace.bytecode_ops) - 1,
        "quarantine_later": bool(trace.quarantined_skill_id) and step_index < len(trace.bytecode_ops) - 1,
        "source_is_success_trace": trace.success,
        "source_is_failure_trace": not trace.success,
        "source_is_adversarial_trace": trace.evaluator_metadata.get("split") == "adversarial",
        "source_is_negative_trace": trace.evaluator_metadata.get("split") == "negative",
    }


def _assign_split(trace: TraceLedgerRecord, trace_index: int) -> str:
    source_split = str(trace.evaluator_metadata.get("split", "train"))
    if source_split == "adversarial":
        return "adversarial"
    if source_split == "negative":
        return "negative"
    if source_split == "heldout":
        variant_id = str(trace.evaluator_metadata.get("variant_id", ""))
        digits = "".join(ch for ch in variant_id if ch.isdigit())
        if digits:
            return "validation" if int(digits) % 2 else "test"
        return "validation" if int(stable_hash(trace.evaluator_metadata), 16) % 2 else "test"
    return "train"


def _gate_source(trace: TraceLedgerRecord) -> str:
    gate = str(trace.evaluator_metadata.get("gate_source", ""))
    if gate in {"gate3", "gate3b", "gate4"}:
        return gate
    split = str(trace.evaluator_metadata.get("split", ""))
    if split in {"adversarial", "negative"} or trace.fallback_used or trace.quarantined_skill_id:
        return "gate3b"
    if split in {"train", "heldout"}:
        return "gate3"
    return "gate4"


def _workspace_summary(recent_summaries: list[str], step_index: int) -> str:
    return f"visible_steps={step_index}; recent={len(recent_summaries)}"


def _previous_tool_summary(visible_tools: list[dict[str, Any]], step_index: int) -> str:
    if not visible_tools:
        return ""
    index = max(0, min(step_index, len(visible_tools) - 1))
    tool = visible_tools[index]
    return f"{tool.get('operation')}:{tool.get('summary')}:{tool.get('success')}"


def _verifier_context_summary(trace: TraceLedgerRecord) -> str:
    passed = sum(1 for output in trace.verifier_outputs if output.get("success"))
    failed = sum(1 for output in trace.verifier_outputs if not output.get("success"))
    return f"verifier_outputs={len(trace.verifier_outputs)}; passed={passed}; failed={failed}"


def _available_skill_cards(trace: TraceLedgerRecord) -> list[dict[str, Any]]:
    if trace.reused_skill_id or any(action == CognitiveOp.RETRIEVE_SKILL.value for action in trace.bytecode_ops):
        return [{"skill_id": trace.reused_skill_id or "visible_skill", "description": "verified visible repair skill", "status": "available"}]
    return []


def _agent_safe_card_dict(card: MemoryCard, evaluator_metadata: dict[str, Any]) -> dict[str, Any]:
    payload = card.to_dict()
    payload["claim"] = _redact_oracle_terms(str(payload.get("claim", "")), evaluator_metadata)
    return payload


def _redact_oracle_terms(text: str, evaluator_metadata: dict[str, Any]) -> str:
    redacted = text
    for term in _forbidden_metadata_values(evaluator_metadata):
        redacted = redacted.replace(term, "visible_repair_pattern")
    return redacted


def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.write_text("".join(canonical_json(row) + "\n" for row in rows), encoding="utf-8")


def _split(transition: ControllerTransition) -> str:
    return str(transition.metadata["split"])


def _read_transitions(path: Path) -> list[ControllerTransition]:
    return [
        ControllerTransition.from_dict(json.loads(line))
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]


def _source_gate_counts(transitions: list[ControllerTransition]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for transition in transitions:
        gate = str(transition.metadata.get("gate_source", "gate4"))
        counts[gate] = counts.get(gate, 0) + 1
    return dict(sorted(counts.items()))


def _forensics_summary(
    transitions: list[ControllerTransition],
    validation: DatasetValidationReport,
    source_trace_count: int,
    split_counts: dict[str, int],
    distinct_action_count: int,
    offline_baselines: dict[str, Any],
) -> dict[str, Any]:
    duplicate_transition_count = _duplicate_count(transition.stable_hash() for transition in transitions)
    trace_splits: dict[str, set[str]] = {}
    for transition in transitions:
        trace_splits.setdefault(str(transition.metadata.get("trace_id", "")), set()).add(_split(transition))
    duplicate_trace_cross_split_count = sum(1 for splits in trace_splits.values() if len(splits) > 1)
    best = dict(offline_baselines.get("best_baseline") or {})
    best_name = str(best.get("name", ""))
    best_metrics = dict(offline_baselines.get(best_name) or {})
    best_splits = dict(best_metrics.get("splits") or {})
    train_acc = float(dict(best_splits.get("train") or {}).get("top1_accuracy", 0.0))
    validation_acc = float(dict(best_splits.get("validation") or {}).get("top1_accuracy", 0.0))
    fallback_transition_count = sum(1 for transition in transitions if transition.labels.get("used_fallback_later"))
    quarantine_transition_count = sum(1 for transition in transitions if transition.labels.get("quarantine_later"))
    leakage_count = validation.oracle_leak_count + validation.canary_leak_count
    return {
        "source_trace_count": source_trace_count,
        "transition_count": len(transitions),
        "train_count": split_counts.get("train", 0),
        "validation_count": split_counts.get("validation", 0),
        "test_count": split_counts.get("test", 0),
        "adversarial_count": split_counts.get("adversarial", 0),
        "negative_count": split_counts.get("negative", 0),
        "fallback_transition_count": fallback_transition_count,
        "quarantine_transition_count": quarantine_transition_count,
        "distinct_action_count": distinct_action_count,
        "split_collision_count": duplicate_trace_cross_split_count,
        "duplicate_transition_count": duplicate_transition_count,
        "duplicate_trace_cross_split_count": duplicate_trace_cross_split_count,
        "leakage_count": leakage_count,
        "canary_leak_count": validation.canary_leak_count,
        "oracle_field_leak_count": validation.oracle_field_leak_count,
        "unsupported_transition_count": validation.unsupported_source_count,
        "illegal_action_count": validation.illegal_action_count,
        "invalid_arg_count": validation.invalid_arg_count,
        "action_mask_violation_count": validation.action_mask_violation_count,
        "unsafe_action_count": validation.unsafe_action_count,
        "overfit_gap_train_vs_validation": max(0.0, train_acc - validation_acc),
    }


def _canary_audit(validation: DatasetValidationReport) -> dict[str, Any]:
    return {
        "canary_terms_tested": validation.canary_terms_tested,
        "canary_leak_count": validation.canary_leak_count,
        "hidden_value_leak_count": validation.hidden_value_leak_count,
        "canary_passed": validation.canary_terms_tested > 0 and validation.canary_leak_count == 0,
    }


def _grounding_context_audit(
    transitions: list[ControllerTransition],
    sequence_shortcut_audit: dict[str, Any],
) -> dict[str, Any]:
    total = len(transitions)
    observed = 0
    workspace = 0
    memory = 0
    verifier = 0
    skill = 0
    tool = 0
    for transition in transitions:
        model_input = transition.model_input
        observation = dict(model_input.get("agent_visible_observation") or {})
        if observation.get("task_description") or observation.get("recent_step_summaries"):
            observed += 1
        if model_input.get("workspace_summary"):
            workspace += 1
        if model_input.get("retrieved_memory_cards"):
            memory += 1
        if model_input.get("verifier_context_summary"):
            verifier += 1
        if model_input.get("available_skill_cards"):
            skill += 1
        if observation.get("visible_tool_summaries"):
            tool += 1
    grounding_required = int(sequence_shortcut_audit.get("grounding_required_transition_count", 0))
    return {
        "transition_count": total,
        "agent_visible_observation_count": observed,
        "workspace_summary_count": workspace,
        "memory_context_transition_count": memory,
        "verifier_context_transition_count": verifier,
        "skill_context_transition_count": skill,
        "tool_context_transition_count": tool,
        "grounding_required_transition_count": grounding_required,
        "grounding_context_present": all(
            count > 0 for count in [observed, workspace, memory, verifier, skill, tool, grounding_required]
        ),
    }


def _shortcut_accuracy(
    train: list[ControllerTransition],
    evaluation: list[ControllerTransition],
    key_fn: object,
) -> float:
    if not train or not evaluation:
        return 0.0
    grouped: dict[str, dict[str, int]] = {}
    action_counts = _action_counts(train)
    fallback = sorted(action_counts, key=lambda item: (-action_counts[item], item))[0]
    for transition in train:
        key = str(key_fn(transition))  # type: ignore[misc]
        action = str(transition.labels["target_action"])
        grouped.setdefault(key, {})[action] = grouped.setdefault(key, {}).get(action, 0) + 1
    table = {
        key: sorted(counts, key=lambda action: (-counts[action], action))[0]
        for key, counts in grouped.items()
    }
    correct = 0
    for transition in evaluation:
        key = str(key_fn(transition))  # type: ignore[misc]
        correct += int(table.get(key, fallback) == str(transition.labels["target_action"]))
    return correct / len(evaluation)


def _ambiguous_key_stats(
    transitions: list[ControllerTransition],
    key_fn: object,
) -> dict[str, int]:
    grouped: dict[str, set[str]] = {}
    counts: dict[str, int] = {}
    for transition in transitions:
        key = str(key_fn(transition))  # type: ignore[misc]
        grouped.setdefault(key, set()).add(str(transition.labels["target_action"]))
        counts[key] = counts.get(key, 0) + 1
    ambiguous = {key for key, actions in grouped.items() if len(actions) > 1}
    return {
        "ambiguous_key_count": len(ambiguous),
        "ambiguous_transition_count": sum(counts[key] for key in ambiguous),
    }


def _duplicate_count(values: object) -> int:
    seen: set[str] = set()
    duplicates = 0
    for value in values:  # type: ignore[union-attr]
        text = str(value)
        if text in seen:
            duplicates += 1
        else:
            seen.add(text)
    return duplicates


def _action_counts(transitions: list[ControllerTransition]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for transition in transitions:
        action = str(transition.labels["target_action"])
        counts[action] = counts.get(action, 0) + 1
    return counts


def _evaluate_policy(policy: object, splits: dict[str, list[ControllerTransition]]) -> dict[str, Any]:
    split_metrics = {split: _evaluate_policy_split(policy, transitions) for split, transitions in splits.items()}
    total = sum(metrics["total"] for metrics in split_metrics.values())
    correct = sum(metrics["correct"] for metrics in split_metrics.values())
    invalid = sum(metrics["invalid_action_count"] for metrics in split_metrics.values())
    return {
        "top1_accuracy": correct / total if total else 0.0,
        "legal_action_rate": (total - invalid) / total if total else 0.0,
        "invalid_action_count": invalid,
        "splits": split_metrics,
        "controls_live_agent": False,
        "leak_count": 0,
        "regression_count": 0,
    }


def _evaluate_policy_split(policy: object, transitions: list[ControllerTransition]) -> dict[str, Any]:
    legal_actions = {op.value for op in CognitiveOp}
    correct = 0
    invalid = 0
    per_action: dict[str, list[int]] = {}
    terminal_total = 0
    terminal_correct = 0
    for transition in transitions:
        prediction = str(policy.predict(transition))
        target = str(transition.labels["target_action"])
        if prediction not in legal_actions:
            invalid += 1
        matched = prediction == target
        correct += int(matched)
        per_action.setdefault(target, []).append(int(matched))
        if target == CognitiveOp.ANSWER.value:
            terminal_total += 1
            terminal_correct += int(matched)
    macro = sum(sum(values) / len(values) for values in per_action.values()) / len(per_action) if per_action else 0.0
    total = len(transitions)
    return {
        "total": total,
        "correct": correct,
        "top1_accuracy": correct / total if total else 0.0,
        "legal_action_rate": (total - invalid) / total if total else 0.0,
        "invalid_action_count": invalid,
        "macro_action_accuracy": macro,
        "terminal_action_accuracy": terminal_correct / terminal_total if terminal_total else 0.0,
    }


def _contains_forbidden_value(model_input: dict[str, Any], metadata: dict[str, Any]) -> bool:
    evaluator = dict(metadata.get("evaluator_metadata") or {})
    return bool(_forbidden_value_hits(model_input, _forbidden_metadata_values(evaluator)))


def _forbidden_value_hits(model_input: dict[str, Any], forbidden_values: list[str]) -> list[str]:
    text = json.dumps(model_input, sort_keys=True)
    return sorted({value for value in forbidden_values if value in text})


def _json_safe(value: object) -> bool:
    try:
        canonical_json(value)
    except (TypeError, ValueError):
        return False
    return True


def _forbidden_metadata_values(evaluator_metadata: dict[str, Any]) -> list[str]:
    keys = [
        "task_family",
        "defect_kind",
        "variant_id",
        "target_file",
        "expected_file_name",
        "expected_symbol",
        "expected_fix",
        "expected_repair",
        "old_text",
        "new_text",
        "good_content",
    ]
    values: list[str] = []
    for key in keys:
        value = evaluator_metadata.get(key)
        if isinstance(value, str) and len(value) >= 4:
            values.append(value)
    return values
