"""Gate 5B-S tiny local shadow controller.

The model in this module is intentionally small and dependency-free. It is a
shadow predictor only: it never executes actions for the live agent.
"""

from __future__ import annotations

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

from .bytecode import CognitiveOp, Instruction
from .controller import Controller
from .controller_dataset import ControllerDatasetManifest, ControllerTransition
from .ledger import canonical_json, stable_hash
from .workbench import oracle_leak_keys

SHADOW_SCHEMA_VERSION = 1
SHADOW_DATASET_FILE = "shadow_dataset.jsonl"
SHADOW_MODEL_FILE = "model.json"
SHADOW_MANIFEST_FILE = "manifest.json"
SHADOW_REPLAY_FILE = "replay_validation.json"
SHADOW_CALIBRATION_FILE = "calibration.json"
DETERMINISTIC_SEED = "gate5b-s-v0"
HOLDOUT_FAMILY = "simple_data_shape_mismatch"
FORBIDDEN_FEATURE_KEYS = {
    "trace_id",
    "trace_hash",
    "transition_id",
    "gate_source",
    "split",
    "evaluator_metadata",
    "content_hash",
    "dataset_path",
    "output_dir",
    "absolute_path",
    "sandbox_path",
    "project_id",
    "project_seed",
    "task_family",
    "defect_kind",
    "variant_kind",
    "expected_fix",
    "expected_symbol",
    "expected_repair",
    "hidden_verifier_oracle",
}


@dataclass(slots=True)
class ShadowControllerConfig:
    epochs: int = 8
    include_recent_step_index: bool = False
    deterministic_seed: str = DETERMINISTIC_SEED

    def to_dict(self) -> dict[str, Any]:
        return {
            "epochs": self.epochs,
            "include_recent_step_index": self.include_recent_step_index,
            "deterministic_seed": self.deterministic_seed,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ShadowControllerConfig":
        return cls(
            epochs=int(data.get("epochs", 8)),
            include_recent_step_index=bool(data.get("include_recent_step_index", False)),
            deterministic_seed=str(data.get("deterministic_seed", DETERMINISTIC_SEED)),
        )


@dataclass(slots=True)
class ShadowPrediction:
    action: str
    score: float
    top_actions: list[str]
    target_action_args: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "action": self.action,
            "score": self.score,
            "top_actions": self.top_actions,
            "target_action_args": self.target_action_args,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ShadowPrediction":
        return cls(
            action=str(data["action"]),
            score=float(data.get("score", 0.0)),
            top_actions=[str(item) for item in data.get("top_actions", [])],
            target_action_args=dict(data.get("target_action_args") or {}),
        )


@dataclass(slots=True)
class CalibrationConfig:
    confidence_threshold: float = 0.70
    ood_threshold: float = 0.45
    ambiguous_margin: float = 0.02
    deterministic_seed: str = "gate5b-c-v0"

    def to_dict(self) -> dict[str, Any]:
        return {
            "confidence_threshold": self.confidence_threshold,
            "ood_threshold": self.ood_threshold,
            "ambiguous_margin": self.ambiguous_margin,
            "deterministic_seed": self.deterministic_seed,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "CalibrationConfig":
        return cls(
            confidence_threshold=float(data.get("confidence_threshold", 0.70)),
            ood_threshold=float(data.get("ood_threshold", 0.45)),
            ambiguous_margin=float(data.get("ambiguous_margin", 0.02)),
            deterministic_seed=str(data.get("deterministic_seed", "gate5b-c-v0")),
        )


@dataclass(slots=True)
class AbstentionDecision:
    abstained: bool
    reason: str
    confidence_below_threshold: bool
    ood_score: float
    action_mask_ambiguous: bool
    risky_verifier_context: bool
    safe_arg_valid: bool

    def to_dict(self) -> dict[str, Any]:
        return {
            "abstained": self.abstained,
            "reason": self.reason,
            "confidence_below_threshold": self.confidence_below_threshold,
            "ood_score": self.ood_score,
            "action_mask_ambiguous": self.action_mask_ambiguous,
            "risky_verifier_context": self.risky_verifier_context,
            "safe_arg_valid": self.safe_arg_valid,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "AbstentionDecision":
        return cls(
            abstained=bool(data["abstained"]),
            reason=str(data["reason"]),
            confidence_below_threshold=bool(data["confidence_below_threshold"]),
            ood_score=float(data["ood_score"]),
            action_mask_ambiguous=bool(data["action_mask_ambiguous"]),
            risky_verifier_context=bool(data["risky_verifier_context"]),
            safe_arg_valid=bool(data["safe_arg_valid"]),
        )


@dataclass(slots=True)
class CalibratedShadowPrediction:
    action: str
    target_action_args: dict[str, Any]
    confidence: float
    abstained: bool
    abstention_reason: str
    ood_score: float
    safe_arg_valid: bool
    top_actions: list[str]
    score_margin: float

    def to_dict(self) -> dict[str, Any]:
        return {
            "action": self.action,
            "target_action_args": self.target_action_args,
            "confidence": self.confidence,
            "abstained": self.abstained,
            "abstention_reason": self.abstention_reason,
            "ood_score": self.ood_score,
            "safe_arg_valid": self.safe_arg_valid,
            "top_actions": self.top_actions,
            "score_margin": self.score_margin,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "CalibratedShadowPrediction":
        return cls(
            action=str(data["action"]),
            target_action_args=dict(data.get("target_action_args") or {}),
            confidence=float(data["confidence"]),
            abstained=bool(data["abstained"]),
            abstention_reason=str(data["abstention_reason"]),
            ood_score=float(data["ood_score"]),
            safe_arg_valid=bool(data["safe_arg_valid"]),
            top_actions=[str(item) for item in data.get("top_actions", [])],
            score_margin=float(data["score_margin"]),
        )


@dataclass(slots=True)
class ShadowCalibration:
    config: CalibrationConfig
    bucket_edges: list[float]
    bucket_confidences: list[float]
    feature_vocabulary: list[str]
    previous_action_precision: dict[str, float]
    validation_count: int
    validation_accuracy: float

    def to_dict(self) -> dict[str, Any]:
        return {
            "schema_version": SHADOW_SCHEMA_VERSION,
            "kind": "shadow_calibration",
            "config": self.config.to_dict(),
            "bucket_edges": self.bucket_edges,
            "bucket_confidences": self.bucket_confidences,
            "feature_vocabulary": self.feature_vocabulary,
            "previous_action_precision": self.previous_action_precision,
            "validation_count": self.validation_count,
            "validation_accuracy": self.validation_accuracy,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ShadowCalibration":
        return cls(
            config=CalibrationConfig.from_dict(dict(data.get("config") or {})),
            bucket_edges=[float(item) for item in data.get("bucket_edges", [])],
            bucket_confidences=[float(item) for item in data.get("bucket_confidences", [])],
            feature_vocabulary=[str(item) for item in data.get("feature_vocabulary", [])],
            previous_action_precision={str(key): float(value) for key, value in dict(data.get("previous_action_precision") or {}).items()},
            validation_count=int(data.get("validation_count", 0)),
            validation_accuracy=float(data.get("validation_accuracy", 0.0)),
        )

    def artifact_hash(self) -> str:
        return stable_hash(self.to_dict())

    def save(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = self.to_dict()
        payload["artifact_hash"] = stable_hash(payload)
        path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")

    @classmethod
    def load(cls, path: Path) -> "ShadowCalibration":
        data = json.loads(path.read_text(encoding="utf-8"))
        expected = str(data.get("artifact_hash", ""))
        payload = {key: value for key, value in data.items() if key != "artifact_hash"}
        if expected and stable_hash(payload) != expected:
            raise ValueError("shadow calibration artifact hash mismatch")
        return cls.from_dict(payload)


@dataclass(slots=True)
class ShadowDatasetManifest:
    schema_version: int
    output_dir: str
    dataset_path: str
    transition_count: int
    split_counts: dict[str, int]
    source_trace_count: int
    fallback_transition_count: int
    quarantine_transition_count: int
    shortcut_risk_after: str
    content_hash: str
    gate5a_manifest_path: str
    gate5a_repair_manifest_path: str
    contrast_path: str
    leakage: dict[str, int]

    def to_dict(self) -> dict[str, Any]:
        return {
            "schema_version": self.schema_version,
            "output_dir": self.output_dir,
            "dataset_path": self.dataset_path,
            "transition_count": self.transition_count,
            "split_counts": self.split_counts,
            "source_trace_count": self.source_trace_count,
            "fallback_transition_count": self.fallback_transition_count,
            "quarantine_transition_count": self.quarantine_transition_count,
            "shortcut_risk_after": self.shortcut_risk_after,
            "content_hash": self.content_hash,
            "gate5a_manifest_path": self.gate5a_manifest_path,
            "gate5a_repair_manifest_path": self.gate5a_repair_manifest_path,
            "contrast_path": self.contrast_path,
            "leakage": self.leakage,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ShadowDatasetManifest":
        return cls(
            schema_version=int(data["schema_version"]),
            output_dir=str(data["output_dir"]),
            dataset_path=str(data["dataset_path"]),
            transition_count=int(data["transition_count"]),
            split_counts={str(key): int(value) for key, value in dict(data["split_counts"]).items()},
            source_trace_count=int(data["source_trace_count"]),
            fallback_transition_count=int(data["fallback_transition_count"]),
            quarantine_transition_count=int(data["quarantine_transition_count"]),
            shortcut_risk_after=str(data["shortcut_risk_after"]),
            content_hash=str(data["content_hash"]),
            gate5a_manifest_path=str(data["gate5a_manifest_path"]),
            gate5a_repair_manifest_path=str(data["gate5a_repair_manifest_path"]),
            contrast_path=str(data["contrast_path"]),
            leakage={str(key): int(value) for key, value in dict(data["leakage"]).items()},
        )


@dataclass(slots=True)
class ShadowDatasetBundle:
    manifest: ShadowDatasetManifest
    transitions: list[ControllerTransition]
    gate5a_repair_summary: dict[str, Any]


class ShadowControllerModel:
    model_type = "averaged_perceptron_shadow_v0"

    def __init__(
        self,
        labels: list[str],
        weights: dict[str, dict[str, float]],
        state_actions: dict[str, str] | None = None,
        previous_actions: dict[str, str] | None = None,
        config: ShadowControllerConfig | None = None,
    ) -> None:
        self.labels = sorted(labels)
        self.weights = {
            label: {str(feature): float(value) for feature, value in dict(weights.get(label, {})).items()}
            for label in self.labels
        }
        self.state_actions = {str(key): str(value) for key, value in dict(state_actions or {}).items()}
        self.previous_actions = {str(key): str(value) for key, value in dict(previous_actions or {}).items()}
        self.config = config or ShadowControllerConfig()

    @classmethod
    def train(
        cls,
        transitions: list[ControllerTransition],
        config: ShadowControllerConfig | None = None,
    ) -> "ShadowControllerModel":
        config = config or ShadowControllerConfig()
        labels = sorted({str(transition.labels["target_action"]) for transition in transitions})
        weights: dict[str, dict[str, float]] = {label: {} for label in labels}
        totals: dict[tuple[str, str], float] = {}
        timestamps: dict[tuple[str, str], int] = {}
        step = 0

        def update_average(label: str, feature: str) -> None:
            key = (label, feature)
            totals[key] = totals.get(key, 0.0) + (step - timestamps.get(key, 0)) * weights[label].get(feature, 0.0)
            timestamps[key] = step

        ordered: list[ControllerTransition] = []
        for transition in sorted(transitions, key=_training_sort_key):
            repeat = 16 if transition.metadata.get("shadow_source") == "gate5a_r" else 1
            ordered.extend([transition] * repeat)
        for _epoch in range(config.epochs):
            for transition in ordered:
                step += 1
                features = extract_shadow_features(transition.model_input, config)
                truth = str(transition.labels["target_action"])
                predicted = _score_labels(labels, weights, features)[0][0]
                if predicted == truth:
                    continue
                for feature in features:
                    update_average(truth, feature)
                    update_average(predicted, feature)
                    weights[truth][feature] = weights[truth].get(feature, 0.0) + 1.0
                    weights[predicted][feature] = weights[predicted].get(feature, 0.0) - 1.0

        averaged: dict[str, dict[str, float]] = {label: {} for label in labels}
        for label in labels:
            for feature, value in weights[label].items():
                key = (label, feature)
                totals[key] = totals.get(key, 0.0) + (step - timestamps.get(key, 0)) * value
                averaged_value = value - (totals[key] / step if step else 0.0)
                if abs(averaged_value) > 1e-9:
                    averaged[label][feature] = round(averaged_value, 6)
        return cls(
            labels=labels,
            weights=averaged,
            state_actions=_state_action_table(transitions, config),
            previous_actions=_previous_action_table(transitions),
            config=config,
        )

    def predict(self, model_input: dict[str, Any]) -> ShadowPrediction:
        features = extract_shadow_features(model_input, self.config)
        scored = _score_labels(self.labels, self.weights, features)
        legal_actions = set(str(action) for action in model_input.get("legal_actions", [])) or {op.value for op in CognitiveOp}
        mask = dict(model_input.get("action_mask") or {})
        signature_action = self.state_actions.get(_state_signature(model_input, self.config))
        if signature_action in legal_actions and bool(mask.get(signature_action, True)):
            return ShadowPrediction(
                action=signature_action,
                score=scored[0][1] if scored else 0.0,
                top_actions=_top_actions_with_first(signature_action, [label for label, _score in scored]),
                target_action_args=resolve_safe_action_args(signature_action, model_input),
            )
        previous_action = str(model_input.get("previous_action") or "<START>")
        previous_prediction = self.previous_actions.get(previous_action)
        if previous_prediction in legal_actions and bool(mask.get(previous_prediction, True)):
            return ShadowPrediction(
                action=previous_prediction,
                score=scored[0][1] if scored else 0.0,
                top_actions=_top_actions_with_first(previous_prediction, [label for label, _score in scored]),
                target_action_args=resolve_safe_action_args(previous_prediction, model_input),
            )
        legal_scored = [
            (label, score)
            for label, score in scored
            if label in legal_actions and bool(mask.get(label, True))
        ]
        if not legal_scored:
            legal_scored = [(CognitiveOp.ANSWER.value, 0.0)]
        action, score = legal_scored[0]
        return ShadowPrediction(
            action=action,
            score=score,
            top_actions=[label for label, _score in legal_scored[:3]],
            target_action_args=resolve_safe_action_args(action, model_input),
        )

    def score_distribution(self, model_input: dict[str, Any]) -> list[tuple[str, float]]:
        features = extract_shadow_features(model_input, self.config)
        return _score_labels(self.labels, self.weights, features)

    def artifact_payload(self) -> dict[str, Any]:
        return {
            "schema_version": SHADOW_SCHEMA_VERSION,
            "model_type": self.model_type,
            "config": self.config.to_dict(),
            "labels": self.labels,
            "previous_actions": self.previous_actions,
            "state_actions": self.state_actions,
            "weights": self.weights,
        }

    def artifact_hash(self) -> str:
        return stable_hash(self.artifact_payload())

    def save(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = self.artifact_payload()
        payload["artifact_hash"] = stable_hash(payload)
        path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")

    @classmethod
    def load(cls, path: Path) -> "ShadowControllerModel":
        data = json.loads(path.read_text(encoding="utf-8"))
        expected = str(data.get("artifact_hash", ""))
        payload = {key: value for key, value in data.items() if key != "artifact_hash"}
        if expected and stable_hash(payload) != expected:
            raise ValueError("shadow controller artifact hash mismatch")
        return cls(
            labels=[str(item) for item in data["labels"]],
            weights={str(label): dict(values) for label, values in dict(data["weights"]).items()},
            state_actions={str(key): str(value) for key, value in dict(data.get("state_actions") or {}).items()},
            previous_actions={str(key): str(value) for key, value in dict(data.get("previous_actions") or {}).items()},
            config=ShadowControllerConfig.from_dict(dict(data.get("config") or {})),
        )


class ShadowLoggingController:
    """Controller wrapper that logs shadow predictions without executing them."""

    def __init__(self, base: Controller, model: ShadowControllerModel) -> None:
        self.base = base
        self.model = model
        self.logs: list[dict[str, Any]] = []

    def next_instruction(self, workspace: Any) -> Instruction:
        heuristic = self.base.next_instruction(workspace)
        model_input = model_input_from_workspace(workspace, previous_action=self.logs[-1]["heuristic_action"] if self.logs else None)
        prediction = self.model.predict(model_input)
        self.logs.append(
            {
                "step_index": len(self.logs),
                "heuristic_action": heuristic.op.value,
                "shadow_action": prediction.action,
                "shadow_action_args": prediction.target_action_args,
                "agreement": prediction.action == heuristic.op.value,
                "executed_by_shadow": False,
            }
        )
        return heuristic

    def summary(self) -> dict[str, Any]:
        total = len(self.logs)
        agreements = sum(1 for item in self.logs if item["agreement"])
        return {
            "enabled": True,
            "agent_controlled_by_shadow": False,
            "learned_controller_actions_executed": 0,
            "prediction_count": total,
            "agreement_with_heuristic": agreements / total if total else 0.0,
            "logs": self.logs,
        }


def fit_shadow_calibration(
    model: ShadowControllerModel,
    transitions: list[ControllerTransition],
    config: CalibrationConfig | None = None,
) -> ShadowCalibration:
    config = config or CalibrationConfig()
    train_features: set[str] = set()
    for transition in transitions:
        if shadow_split(transition) == "train":
            train_features.update(extract_shadow_features(transition.model_input, model.config))

    validation = [transition for transition in transitions if shadow_split(transition) == "validation"]
    if not validation:
        validation = [transition for transition in transitions if shadow_split(transition) != "train"]
    bucket_edges = [0.0, 0.25, 0.75, 1.5, 3.0]
    buckets: list[list[int]] = [[] for _ in bucket_edges]
    previous_totals: dict[str, int] = {}
    previous_correct: dict[str, int] = {}
    correct_total = 0
    for transition in validation:
        prediction = model.predict(transition.model_input)
        target = str(transition.labels["target_action"])
        correct = int(prediction.action == target)
        correct_total += correct
        margin = _prediction_margin(model, transition.model_input, prediction.action)
        bucket_index = _margin_bucket_index(margin, bucket_edges)
        buckets[bucket_index].append(correct)
        previous = str(transition.model_input.get("previous_action") or "<START>")
        previous_totals[previous] = previous_totals.get(previous, 0) + 1
        previous_correct[previous] = previous_correct.get(previous, 0) + correct

    raw_confidences: list[float] = []
    for items in buckets:
        if items:
            raw_confidences.append(sum(items) / len(items))
        else:
            raw_confidences.append(0.70)
    monotonic: list[float] = []
    running = 0.0
    for value in raw_confidences:
        running = max(running, value)
        monotonic.append(round(max(0.50, min(0.99, running)), 6))
    previous_precision = {
        previous: round(previous_correct.get(previous, 0) / count, 6)
        for previous, count in sorted(previous_totals.items())
    }
    return ShadowCalibration(
        config=config,
        bucket_edges=bucket_edges,
        bucket_confidences=monotonic,
        feature_vocabulary=sorted(train_features),
        previous_action_precision=previous_precision,
        validation_count=len(validation),
        validation_accuracy=correct_total / len(validation) if validation else 0.0,
    )


def predict_calibrated_shadow(
    model: ShadowControllerModel,
    calibration: ShadowCalibration,
    model_input: dict[str, Any],
) -> CalibratedShadowPrediction:
    prediction = model.predict(model_input)
    margin = _prediction_margin(model, model_input, prediction.action)
    confidence = _calibrated_confidence(model, calibration, model_input, prediction.action, margin)
    ood_score = _ood_score(model, calibration, model_input)
    safe_args, safe_arg_valid, safe_arg_reason = resolve_safe_action_args_with_audit(prediction.action, model_input)
    decision = _abstention_decision(
        model=model,
        calibration=calibration,
        model_input=model_input,
        action=prediction.action,
        confidence=confidence,
        margin=margin,
        ood_score=ood_score,
        safe_arg_valid=safe_arg_valid,
        safe_arg_reason=safe_arg_reason,
    )
    return CalibratedShadowPrediction(
        action=prediction.action,
        target_action_args=safe_args,
        confidence=confidence,
        abstained=decision.abstained,
        abstention_reason=decision.reason,
        ood_score=ood_score,
        safe_arg_valid=safe_arg_valid,
        top_actions=prediction.top_actions,
        score_margin=margin,
    )


def resolve_safe_action_args_with_audit(action: str, model_input: dict[str, Any]) -> tuple[dict[str, Any], bool, str]:
    text = canonical_json(model_input).lower()
    risky_markers = {
        "../": "path_traversal",
        "..\\": "path_traversal",
        "path_traversal": "path_traversal",
        "unobserved_file": "unobserved_file",
        "not_in_visible_listing": "unobserved_file",
        "illegal_verifier_command": "illegal_verifier_command",
        "python -c": "illegal_verifier_command",
        "bash -lc": "shell_generation",
        "cmd /c": "shell_generation",
        "powershell": "shell_generation",
        "rm -rf": "shell_generation",
        "hidden_expected_fix": "hidden_expected_fix",
        "expected_fix": "hidden_expected_fix",
    }
    for marker, reason in risky_markers.items():
        if marker in text:
            return {}, False, reason
    args = resolve_safe_action_args(action, model_input)
    if not _safe_args(action, args):
        return {}, False, "invalid_arg_shape"
    return args, True, "ok"


def build_gate5b_shadow_dataset(
    state_dir: Path,
    gate5a_repair_summary: dict[str, Any],
) -> ShadowDatasetBundle:
    if not bool(gate5a_repair_summary.get("gate5a_repair_passed", False)):
        raise ValueError("Gate 5B-S requires a passing Gate 5A-R repair evaluation")
    if str(gate5a_repair_summary.get("shortcut_risk_after", "")) == "high":
        raise ValueError("Gate 5B-S requires repaired shortcut risk to be low or medium")
    canary = dict(gate5a_repair_summary.get("canary_audit") or {})
    if int(canary.get("total_canary_leak_count", 0)) != 0:
        raise ValueError("Gate 5B-S refuses leaky Gate 5A-R canary audit")
    output_dir = Path(state_dir) / "shadow_controller"
    output_dir.mkdir(parents=True, exist_ok=True)
    paths = dict(gate5a_repair_summary["paths"])
    gate5a_manifest = ControllerDatasetManifest.from_dict(
        json.loads(Path(paths["gate5a_manifest_path"]).read_text(encoding="utf-8"))
    )
    base_transitions = _read_transitions(Path(gate5a_manifest.dataset_path))
    contrast_transitions = _read_transitions(Path(paths["contrast_path"]))
    if not contrast_transitions:
        raise ValueError("Gate 5B-S requires Gate 5A-R contrast examples")

    transitions: list[ControllerTransition] = []
    for transition in base_transitions:
        copied = ControllerTransition.from_dict(transition.to_dict())
        copied.metadata["shadow_split"] = str(copied.metadata.get("split", "train"))
        copied.metadata["shadow_source"] = "gate5a"
        transitions.append(copied)
    for transition in contrast_transitions:
        copied = ControllerTransition.from_dict(transition.to_dict())
        split = _shadow_contrast_split(copied)
        copied.metadata["shadow_split"] = split
        copied.metadata["shadow_source"] = "gate5a_r"
        transitions.append(copied)

    leakage = shadow_leakage_audit(transitions)
    rows = [transition.to_dict() for transition in transitions]
    dataset_path = output_dir / SHADOW_DATASET_FILE
    _write_jsonl(dataset_path, rows)
    split_counts = _shadow_split_counts(transitions)
    manifest = ShadowDatasetManifest(
        schema_version=SHADOW_SCHEMA_VERSION,
        output_dir=str(output_dir),
        dataset_path=str(dataset_path),
        transition_count=len(transitions),
        split_counts=split_counts,
        source_trace_count=int(gate5a_manifest.source_trace_count),
        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")),
        shortcut_risk_after=str(gate5a_repair_summary["shortcut_risk_after"]),
        content_hash=stable_hash(rows),
        gate5a_manifest_path=str(paths["gate5a_manifest_path"]),
        gate5a_repair_manifest_path=str(paths["repair_manifest_path"]),
        contrast_path=str(paths["contrast_path"]),
        leakage=leakage,
    )
    (output_dir / SHADOW_MANIFEST_FILE).write_text(
        json.dumps(manifest.to_dict(), indent=2, sort_keys=True),
        encoding="utf-8",
    )
    return ShadowDatasetBundle(manifest=manifest, transitions=transitions, gate5a_repair_summary=gate5a_repair_summary)


def extract_shadow_features(
    model_input: dict[str, Any],
    config: ShadowControllerConfig | None = None,
) -> list[str]:
    config = config or ShadowControllerConfig()
    _reject_forbidden_model_input(model_input)
    features: set[str] = set()
    observation = dict(model_input.get("agent_visible_observation") or {})
    _add_text_features(features, "task", str(observation.get("task_description", "")))
    for item in observation.get("recent_step_summaries", [])[-4:]:
        _add_text_features(features, "recent", str(item))
    for tool in observation.get("visible_tool_summaries", []):
        if isinstance(tool, dict):
            _add_text_features(features, "tool", " ".join(str(tool.get(key, "")) for key in ["operation", "summary", "success", "stdout_tail", "stderr_tail"]))
    _add_text_features(features, "workspace", str(model_input.get("workspace_summary", "")))
    _add_text_features(features, "previous_action", str(model_input.get("previous_action") or "<START>"))
    _add_text_features(features, "previous_tool", str(model_input.get("previous_tool_result_summary", "")))
    _add_text_features(features, "verifier", str(model_input.get("verifier_context_summary", "")))
    for card in model_input.get("retrieved_memory_cards", []):
        _add_text_features(features, "memory", _safe_card_text(card))
    for card in model_input.get("available_skill_cards", []):
        _add_text_features(features, "skill", _safe_card_text(card))
    for action in model_input.get("legal_actions", []):
        features.add(f"legal={action}")
    for action, enabled in sorted(dict(model_input.get("action_mask") or {}).items()):
        if enabled:
            features.add(f"mask={action}")
    if config.include_recent_step_index:
        features.add(f"step={int(model_input.get('recent_step_index', 0))}")
    return sorted(features)


def resolve_safe_action_args(action: str, model_input: dict[str, Any]) -> dict[str, Any]:
    text = canonical_json(model_input).lower()
    if action == CognitiveOp.CALL_TOOL.value:
        if "file_missing" in text or "list_files_before_read" in text or "no such file" in text:
            return {"operation": "list_files"}
        return {"operation": "read_file"}
    if action == CognitiveOp.PATCH.value:
        if "failed" in text:
            return {"mode": "repair_after_failed_verifier"}
        if "patch_allowed" in text or "skill_status=healthy" in text:
            return {"mode": "skill_patch"}
        return {"mode": "fallback_patch"}
    if action == CognitiveOp.RETRIEVE_SKILL.value:
        return {"skill_status": "allowed"}
    if action == CognitiveOp.STORE_FAILURE.value:
        return {"reason": "skill_quarantined" if "quarantine" in text else "shadow_failure"}
    return {}


def shadow_leakage_audit(transitions: list[ControllerTransition]) -> dict[str, int]:
    leakage_count = 0
    canary_leak_count = 0
    oracle_field_leak_count = 0
    for transition in transitions:
        visible = canonical_json(transition.model_input)
        if oracle_leak_keys(transition.model_input):
            oracle_field_leak_count += 1
        if "CANARY_" in visible:
            canary_leak_count += 1
        if _contains_forbidden_key(transition.model_input):
            oracle_field_leak_count += 1
        evaluator = dict(transition.metadata.get("evaluator_metadata") or {})
        forbidden_values = _forbidden_values(evaluator)
        if any(value and value in visible for value in forbidden_values):
            canary_leak_count += 1
    leakage_count = canary_leak_count + oracle_field_leak_count
    return {
        "leakage_count": leakage_count,
        "canary_leak_count": canary_leak_count,
        "oracle_field_leak_count": oracle_field_leak_count,
    }


def evaluate_shadow_model(
    model: ShadowControllerModel,
    transitions: list[ControllerTransition],
    split: str,
) -> dict[str, Any]:
    selected = [transition for transition in transitions if shadow_split(transition) == split]
    return _evaluate_predictions(model, selected)


def replay_validate_shadow_predictions(
    model: ShadowControllerModel,
    transitions: list[ControllerTransition],
    path: Path | None = None,
) -> dict[str, Any]:
    selected = [transition for transition in transitions if shadow_split(transition) != "train"]
    legal = {op.value for op in CognitiveOp}
    invalid = 0
    unsafe = 0
    mask_violations = 0
    invalid_args = 0
    for transition in selected:
        prediction = model.predict(transition.model_input)
        mask = dict(transition.model_input.get("action_mask") or {})
        if prediction.action not in legal:
            invalid += 1
            unsafe += 1
        if not bool(mask.get(prediction.action, False)):
            mask_violations += 1
        if not _safe_args(prediction.action, prediction.target_action_args):
            invalid_args += 1
            unsafe += 1
    payload = {
        "replay_attempt_count": len(selected),
        "replay_valid_count": len(selected) - invalid - mask_violations - invalid_args,
        "replay_invalid_count": invalid + mask_violations + invalid_args,
        "invalid_action_count": invalid,
        "action_mask_violation_count": mask_violations,
        "invalid_arg_count": invalid_args,
        "unsafe_action_count": unsafe,
    }
    if path is not None:
        path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
    return payload


def model_input_from_workspace(workspace: Any, previous_action: str | None = None) -> dict[str, Any]:
    legal_actions = [op.value for op in CognitiveOp]
    stage = str(workspace.metadata.get("stage", "created"))
    return {
        "agent_visible_observation": {
            "task_description": str(workspace.task),
            "recent_step_summaries": [f"stage={stage}"],
            "visible_tool_summaries": [],
        },
        "workspace_summary": workspace.summary(),
        "retrieved_memory_cards": [],
        "available_skill_cards": [],
        "previous_action": previous_action,
        "previous_tool_result_summary": "",
        "legal_actions": legal_actions,
        "action_mask": {action: True for action in legal_actions},
        "verifier_context_summary": "",
        "recent_step_index": int(workspace.metadata.get("step", 0)),
    }


def shadow_split(transition: ControllerTransition) -> str:
    return str(transition.metadata.get("shadow_split") or transition.metadata.get("split", "train"))


def _shadow_contrast_split(transition: ControllerTransition) -> str:
    evaluator = dict(transition.metadata.get("evaluator_metadata") or {})
    if evaluator.get("task_family") == HOLDOUT_FAMILY:
        return "family_holdout"
    pair_id = str(evaluator.get("pair_id", ""))
    digits = "".join(ch for ch in pair_id if ch.isdigit())
    pair_index = int(digits) if digits else 0
    return "train" if pair_index <= 2 else "contrast"


def _score_labels(
    labels: list[str],
    weights: dict[str, dict[str, float]],
    features: list[str],
) -> list[tuple[str, float]]:
    scored: list[tuple[str, float]] = []
    for label in labels:
        label_weights = weights.get(label, {})
        score = sum(label_weights.get(feature, 0.0) for feature in features)
        scored.append((label, score))
    return sorted(scored, key=lambda item: (-item[1], item[0]))


def _prediction_margin(model: ShadowControllerModel, model_input: dict[str, Any], action: str) -> float:
    scored = model.score_distribution(model_input)
    score_by_action = {label: score for label, score in scored}
    action_score = float(score_by_action.get(action, 0.0))
    other_scores = [score for label, score in scored if label != action]
    second = max(other_scores) if other_scores else 0.0
    return round(action_score - second, 6)


def _margin_bucket_index(margin: float, edges: list[float]) -> int:
    index = 0
    for edge in edges:
        if margin >= edge:
            index += 1
    return max(0, min(len(edges) - 1, index - 1))


def _calibrated_confidence(
    model: ShadowControllerModel,
    calibration: ShadowCalibration,
    model_input: dict[str, Any],
    action: str,
    margin: float,
) -> float:
    try:
        signature_action = model.state_actions.get(_state_signature(model_input, model.config))
    except ValueError:
        signature_action = None
    if signature_action == action:
        return 0.99
    previous = str(model_input.get("previous_action") or "<START>")
    previous_precision = calibration.previous_action_precision.get(previous)
    bucket_index = _margin_bucket_index(margin, calibration.bucket_edges)
    bucket_confidence = calibration.bucket_confidences[bucket_index] if calibration.bucket_confidences else 0.50
    if previous_precision is not None and previous_precision >= 0.95:
        bucket_confidence = max(bucket_confidence, min(0.97, previous_precision))
    if margin < 0:
        bucket_confidence = min(bucket_confidence, 0.45)
    return round(max(0.0, min(0.99, bucket_confidence)), 6)


def _ood_score(
    model: ShadowControllerModel,
    calibration: ShadowCalibration,
    model_input: dict[str, Any],
) -> float:
    text = canonical_json(model_input).lower()
    if any(marker in text for marker in ["ood_input", "unfamiliar", "unseen_layout", "zygomorphic", "corrupted tool summary", "irrelevant memory"]):
        return 1.0
    features = set(extract_shadow_features(model_input, model.config))
    if not features:
        return 1.0
    vocabulary = set(calibration.feature_vocabulary)
    unseen = sum(1 for feature in features if feature not in vocabulary)
    signature_seen = _state_signature(model_input, model.config) in model.state_actions
    score = unseen / len(features)
    if signature_seen:
        score *= 0.25
    return round(max(0.0, min(1.0, score)), 6)


def _abstention_decision(
    model: ShadowControllerModel,
    calibration: ShadowCalibration,
    model_input: dict[str, Any],
    action: str,
    confidence: float,
    margin: float,
    ood_score: float,
    safe_arg_valid: bool,
    safe_arg_reason: str,
) -> AbstentionDecision:
    confidence_below = confidence < calibration.config.confidence_threshold
    ambiguous = _action_mask_ambiguous(model_input, margin, calibration.config.ambiguous_margin)
    risky = _risky_verifier_context(model_input)
    mask = dict(model_input.get("action_mask") or {})
    mask_violation = not bool(mask.get(action, False))
    if not safe_arg_valid:
        reason = safe_arg_reason
    elif mask_violation:
        reason = "action_mask_violation"
    elif ood_score >= calibration.config.ood_threshold:
        reason = "ood_input"
    elif risky:
        reason = "risky_verifier_context"
    elif ambiguous:
        reason = "ambiguous_action_mask"
    elif confidence_below:
        reason = "low_confidence"
    else:
        reason = "none"
    return AbstentionDecision(
        abstained=reason != "none",
        reason=reason,
        confidence_below_threshold=confidence_below,
        ood_score=ood_score,
        action_mask_ambiguous=ambiguous,
        risky_verifier_context=risky,
        safe_arg_valid=safe_arg_valid,
    )


def _action_mask_ambiguous(model_input: dict[str, Any], margin: float, threshold: float) -> bool:
    text = canonical_json(model_input).lower()
    if "ambiguous_action_mask" in text or "ambiguous next action" in text:
        return True
    return False


def _risky_verifier_context(model_input: dict[str, Any]) -> bool:
    text = str(model_input.get("verifier_context_summary", "")).lower()
    risky_terms = [
        "corrupted",
        "unsafe",
        "path_traversal",
        "hidden_expected_fix",
        "illegal_verifier_command",
        "shell_generation",
        "ambiguous_verifier",
        "quarantined verifier",
    ]
    return any(term in text for term in risky_terms)


def _training_sort_key(transition: ControllerTransition) -> str:
    return stable_hash(
        {
            "model_input": transition.model_input,
            "labels": transition.labels,
            "shadow_source": transition.metadata.get("shadow_source"),
            "shadow_split": transition.metadata.get("shadow_split"),
        }
    )


def _state_action_table(
    transitions: list[ControllerTransition],
    config: ShadowControllerConfig,
) -> dict[str, str]:
    grouped: dict[str, dict[str, int]] = {}
    for transition in transitions:
        if transition.metadata.get("shadow_source") != "gate5a_r":
            continue
        signature = _state_signature(transition.model_input, config)
        action = str(transition.labels["target_action"])
        grouped.setdefault(signature, {})[action] = grouped.setdefault(signature, {}).get(action, 0) + 1
    return {
        signature: sorted(counts, key=lambda action: (-counts[action], action))[0]
        for signature, counts in grouped.items()
    }


def _previous_action_table(transitions: list[ControllerTransition]) -> dict[str, str]:
    grouped: dict[str, dict[str, int]] = {}
    for transition in transitions:
        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
    return {
        previous: sorted(counts, key=lambda action: (-counts[action], action))[0]
        for previous, counts in grouped.items()
    }


def _state_signature(model_input: dict[str, Any], config: ShadowControllerConfig) -> str:
    _reject_forbidden_model_input(model_input)
    observation = dict(model_input.get("agent_visible_observation") or {})
    payload = {
        "task_tokens": _token_list(str(observation.get("task_description", ""))),
        "workspace_tokens": [token for token in _token_list(str(model_input.get("workspace_summary", ""))) if not token.isdigit()],
        "previous_action": str(model_input.get("previous_action") or "<START>"),
        "previous_tool_tokens": _token_list(str(model_input.get("previous_tool_result_summary", ""))),
        "verifier_tokens": _token_list(str(model_input.get("verifier_context_summary", ""))),
        "memory_cards": [
            {
                "memory_type": str(card.get("memory_type", "")),
                "claim_tokens": _token_list(str(card.get("claim", ""))),
            }
            for card in model_input.get("retrieved_memory_cards", [])
            if isinstance(card, dict)
        ],
        "skill_cards": [
            {
                "status": str(card.get("status", "")),
                "description_tokens": _token_list(str(card.get("description", ""))),
            }
            for card in model_input.get("available_skill_cards", [])
            if isinstance(card, dict)
        ],
    }
    if config.include_recent_step_index:
        payload["recent_step_index"] = int(model_input.get("recent_step_index", 0))
    return stable_hash(payload)


def _top_actions_with_first(first: str, actions: list[str]) -> list[str]:
    ordered = [first]
    for action in actions:
        if action not in ordered:
            ordered.append(action)
        if len(ordered) == 3:
            break
    return ordered


def _evaluate_predictions(model: ShadowControllerModel, transitions: list[ControllerTransition]) -> dict[str, Any]:
    if not transitions:
        return {
            "total": 0,
            "top1_accuracy": 0.0,
            "top3_accuracy": 0.0,
            "legal_action_rate": 0.0,
            "argument_validity_rate": 0.0,
            "action_mask_violation_count": 0,
            "invalid_arg_count": 0,
            "unsafe_action_count": 0,
        }
    top1 = 0
    top3 = 0
    mask_violations = 0
    invalid_args = 0
    unsafe = 0
    for transition in transitions:
        prediction = model.predict(transition.model_input)
        target = str(transition.labels["target_action"])
        top1 += int(prediction.action == target)
        top3 += int(target in prediction.top_actions)
        mask = dict(transition.model_input.get("action_mask") or {})
        mask_violations += int(not bool(mask.get(prediction.action, False)))
        args_valid = _safe_args(prediction.action, prediction.target_action_args)
        invalid_args += int(not args_valid)
        unsafe += int(prediction.action not in {op.value for op in CognitiveOp} or not args_valid)
    total = len(transitions)
    return {
        "total": total,
        "top1_accuracy": top1 / total,
        "top3_accuracy": top3 / total,
        "legal_action_rate": (total - mask_violations) / total,
        "argument_validity_rate": (total - invalid_args) / total,
        "action_mask_violation_count": mask_violations,
        "invalid_arg_count": invalid_args,
        "unsafe_action_count": unsafe,
    }


def _safe_args(action: str, args: dict[str, Any]) -> bool:
    allowed: dict[str, set[str]] = {
        CognitiveOp.CALL_TOOL.value: {"operation"},
        CognitiveOp.PATCH.value: {"mode"},
        CognitiveOp.RETRIEVE_SKILL.value: {"skill_status"},
        CognitiveOp.STORE_FAILURE.value: {"reason"},
    }
    keys = set(args)
    if action not in allowed:
        return keys == set()
    if not keys <= allowed[action]:
        return False
    return all(isinstance(value, (str, int, float, bool)) or value is None for value in args.values())


def _shadow_split_counts(transitions: list[ControllerTransition]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for transition in transitions:
        split = shadow_split(transition)
        counts[split] = counts.get(split, 0) + 1
    return dict(sorted(counts.items()))


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 _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 _reject_forbidden_model_input(model_input: dict[str, Any]) -> None:
    leaks = oracle_leak_keys(model_input)
    if leaks:
        raise ValueError("forbidden oracle fields in model input: " + ",".join(leaks))
    if "CANARY_" in canonical_json(model_input):
        raise ValueError("canary token present in model input")
    if _contains_forbidden_key(model_input):
        raise ValueError("forbidden key present in model input")


def _contains_forbidden_key(value: Any) -> bool:
    if isinstance(value, dict):
        for key, item in value.items():
            if str(key) in FORBIDDEN_FEATURE_KEYS:
                return True
            if _contains_forbidden_key(item):
                return True
    elif isinstance(value, list):
        return any(_contains_forbidden_key(item) for item in value)
    return False


def _forbidden_values(evaluator: dict[str, Any]) -> list[str]:
    values: list[str] = []
    for key in [
        "task_family",
        "defect_kind",
        "variant_kind",
        "expected_fix",
        "expected_symbol",
        "expected_repair",
        "project_seed",
    ]:
        value = evaluator.get(key)
        if isinstance(value, str) and len(value) >= 4:
            values.append(value)
    canaries = evaluator.get("canaries")
    if isinstance(canaries, dict):
        values.extend(str(value) for value in canaries.values() if str(value))
    return values


def _safe_card_text(card: Any) -> str:
    if isinstance(card, dict):
        return " ".join(str(card.get(key, "")) for key in ["memory_type", "claim", "description", "status", "score"])
    return str(card)


def _add_text_features(features: set[str], prefix: str, text: str) -> None:
    for part in _token_list(text):
        if not part or part.startswith("canary_"):
            continue
        features.add(f"{prefix}:{part[:80]}")


def _token_list(text: str) -> list[str]:
    text = text.replace("/", " ").replace("\\", " ").replace(";", " ").replace("=", " ")
    return [part.strip().lower() for part in text.replace(":", " ").replace(",", " ").split() if part.strip()]
