environment.py

"""
Citadel — Core Environment (two-agent council step loop)

OpenEnv-compatible. Preserves Bastion v1's single-action step signature
but extends behavior: the submitted IncidentAction is treated as a
Commander PROPOSAL and routed through an Oversight policy before being
applied. In inference, both Commander and Oversight are LLMs; the caller
supplies the Oversight's structured critique via the `oversight_action`
kwarg. When no oversight is provided, the env uses a rule-based default
policy so the env still runs solo for validation / smoke tests.

Per step the env:
  1. Takes the submitted Commander action (with justification + cited lessons).
  2. Invokes the Oversight policy (or uses kwarg-supplied OversightAction).
  3. Routes:
       APPROVE → apply action to the world.
       REVISE  → apply the action (the inference script already handled
                 the revision dialogue and submitted the final action).
       VETO    → do NOT apply; the turn advances with a no-op fallback.
       FLAG    → apply the action; mark as audit-flagged.
  4. Runs the adversary turn for the current generation.
  5. Updates trust scores based on the council outcome.
  6. Records the ProposalRecord in CouncilState.
  7. Runs periodic governance checks (GDPR clock, P1 window, etc.).
  8. Writes the Oversight's post-mortem lesson to the shared playbook.
  9. Computes Commander's step reward + penalties.
 10. Returns the Commander's next observation.
"""

from __future__ import annotations

import random
from typing import Any, Callable, Dict, List, Optional
from uuid import uuid4

from openenv.core.env_server import Environment

from models import (
    IncidentAction,
    IncidentObservation,
    IncidentState,
    CommanderProposal,
    OversightAction,
    OversightDecision,
    ActionType,
    ACTION_NAMES,
    BASTION_V1_ACTIONS,
    NUM_ACTIONS,
    ProposalRecord,
    make_observation,
    is_bastion_action,
    is_governance_action,
    SYSTEM_NAMES,
    SERVICE_SYSTEMS,
)
from dynamics import (
    apply_action as apply_bastion_action,
    tick_pending_recompromise,
    generate_team_messages,
    generate_forensic_report,
)
from adversary import adversary_turn, describe_generation
from adversary_llm import gen4_adversary_turn, make_adversary_client_from_env
from governance import (
    apply_governance_action,
    check_prerequisites,
    record_prereq_violations,
    periodic_governance_check,
    DESTRUCTIVE_BASTION_ACTIONS,
    DATA_HOLDING_SYSTEMS,
)
from trust import (
    update_trust_c2o,
    update_trust_o2c,
    drift_toward_mean,
)
from playbook import Playbook, get_playbook, make_context_tags
from stakeholder_events import (
    roll_new_events,
    expire_overdue_asks,
    try_respond as stakeholder_try_respond,
    asks_as_team_messages,
)
from investor_agent import InvestorAgent, PERSONAS
from reward import (
    compute_commander_step_reward,
    compute_oversight_step_reward,
    compute_penalties,
    compute_final_score,
    compute_bastion_v1_final_score,
    compute_baseline_comparison,
)
from baseline import run_baseline
from tasks import get_task, TaskConfig


# ---------------------------------------------------------------------------
# Default rule-based Oversight policy (used when no LLM oversight is supplied)
# ---------------------------------------------------------------------------

def default_oversight_policy(
    proposal: CommanderProposal,
    state: IncidentState,
    action_history: List[Dict[str, Any]],
) -> OversightAction:
    """Heuristic baseline Oversight — keeps the env self-contained."""
    action_idx = proposal.action
    weakness = ""
    missing_evidence: List[str] = []
    risk = 1
    decision = OversightDecision.APPROVE

    target_name = (
        SYSTEM_NAMES[proposal.target_system]
        if 0 <= proposal.target_system < len(SYSTEM_NAMES) else ""
    )

    if action_idx in DESTRUCTIVE_BASTION_ACTIONS and target_name in DATA_HOLDING_SYSTEMS:
        try:
            sys = state.get_system(target_name)
        except Exception:
            sys = None
        if sys and not sys.investigated:
            decision = OversightDecision.VETO
            risk = 4
            weakness = f"proposal targets uninvestigated data system {target_name}"
            missing_evidence.append("no investigation of target system")

    if decision == OversightDecision.APPROVE and len(proposal.justification.strip()) < 12:
        decision = OversightDecision.REVISE
        risk = 2
        weakness = "justification is too short to assess"
        missing_evidence.append("missing evidence pointer in justification")

    return OversightAction(
        decision=int(decision),
        risk_tier=risk,
        weakness=weakness,
        missing_evidence=missing_evidence,
        counter_proposal=None,
        lesson_text="",
        lesson_tags=[],
    )


OversightPolicy = Callable[
    [CommanderProposal, IncidentState, List[Dict[str, Any]]],
    OversightAction,
]


# ---------------------------------------------------------------------------
# Policy-check helpers for the Oversight observation
# ---------------------------------------------------------------------------

def compute_policy_checks(
    state: IncidentState,
    proposal: CommanderProposal,
) -> Dict[str, Any]:
    """Deterministic red-flags for the Oversight view."""
    action_idx = proposal.action
    target_name = (
        SYSTEM_NAMES[proposal.target_system]
        if 0 <= proposal.target_system < len(SYSTEM_NAMES) else ""
    )
    sys = None
    try:
        sys = state.get_system(target_name)
    except Exception:
        pass

    alert_confidence = 0.0
    for a in reversed(state.alerts):
        if a.source_system == target_name and a.confidence > 0:
            alert_confidence = a.confidence
            break

    prereq_problems = (
        check_prerequisites(state.governance_state, action_idx, target_name, state.hour)
        if is_bastion_action(action_idx) else []
    )

    return {
        "action_name": ACTION_NAMES.get(action_idx, f"unknown_{action_idx}"),
        "target_system": target_name,
        "action_targets_uninvestigated_system": bool(sys and not sys.investigated),
        "action_irreversible": action_idx in DESTRUCTIVE_BASTION_ACTIONS,
        "prerequisites_missing": prereq_problems,
        "alert_confidence_for_target": round(alert_confidence, 2),
        "justification_length": len(proposal.justification),
        "has_cited_lessons": bool(proposal.cited_lessons),
    }


# ---------------------------------------------------------------------------
# CitadelEnvironment
# ---------------------------------------------------------------------------

class CitadelEnvironment(Environment[IncidentAction, IncidentObservation, IncidentState]):
    """
    Two-agent council environment.

    Step signature (OpenEnv-compatible):
        env.step(action: IncidentAction, oversight_action: Optional[OversightAction] = None, ...)

    When oversight_action is provided, it's applied verbatim. When it's None,
    the env falls back to a rule-based default Oversight policy.
    """

    VETO_BUDGET = 4
    FLAG_BUDGET = 2

    def __init__(
        self,
        oversight_policy: Optional[OversightPolicy] = None,
        # --- ablation / feature flags ---
        disable_playbook: bool = False,
        disable_trust_dynamics: bool = False,
        disable_governance_prereqs: bool = False,
        disable_stakeholder_events: bool = False,
        force_adversary_gen: Optional[int] = None,
        adversary_llm_client: Optional[Any] = None,
        # LLM client for investor agent — same OpenAI-compatible client as Commander
        investor_llm_client: Optional[Any] = None,
        investor_model_name: str = "",
    ) -> None:
        super().__init__()
        self._task: Optional[TaskConfig] = None
        self._state: IncidentState = IncidentState()
        self._rng: random.Random = random.Random(42)
        self._commander_action_history: List[Dict[str, Any]] = []
        self._cumulative_commander_reward: float = 0.0
        self._cumulative_oversight_reward: float = 0.0
        self._baseline_state: Optional[IncidentState] = None
        self._done: bool = False
        self._initialized: bool = False
        self._alerts_accurate: bool = False
        self._veto_budget_remaining: int = self.VETO_BUDGET
        self._flag_budget_remaining: int = self.FLAG_BUDGET
        self._oversight_policy: OversightPolicy = (
            oversight_policy or default_oversight_policy
        )
        self._playbook: Playbook = get_playbook()
        self._last_critique: Dict[str, Any] = {}

        # Feature flags — configurable per env instance or via reset kwargs.
        # Used by the ablation harness to disable one layer at a time and
        # measure each layer's contribution independently.
        self.disable_playbook = disable_playbook
        self.disable_trust_dynamics = disable_trust_dynamics
        self.disable_governance_prereqs = disable_governance_prereqs
        self.disable_stakeholder_events = disable_stakeholder_events
        self.force_adversary_gen = force_adversary_gen
        self.adversary_llm_client = adversary_llm_client
        self.investor_llm_client = investor_llm_client
        self.investor_model_name = investor_model_name

        # Investor agent — created once per env, reset each episode
        self._investor_agent: InvestorAgent = InvestorAgent(
            rng=self._rng,
            llm_client=investor_llm_client,
            model_name=investor_model_name,
        )

    # --- reset ------------------------------------------------------------

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> IncidentObservation:
        task_id = kwargs.get("task_id", "easy_1")
        # Per-reset overrides also supported for the ablation harness
        for flag in ("disable_playbook", "disable_trust_dynamics",
                     "disable_governance_prereqs", "disable_stakeholder_events"):
            if flag in kwargs:
                setattr(self, flag, bool(kwargs[flag]))
        adversary_gen = (
            kwargs.get("adversary_gen")
            or self.force_adversary_gen
            or None
        )

        self._task = get_task(task_id)
        self._state = self._task.initial_state.clone()
        self._state.episode_id = episode_id or str(uuid4())
        self._state.step_count = 0
        self._state.task_id = task_id
        self._state.adversary_gen = int(adversary_gen or self._task.default_adversary_gen)

        effective_seed = seed if seed is not None else self._task.seed
        self._rng = random.Random(effective_seed)

        self._commander_action_history = []
        self._cumulative_commander_reward = 0.0
        self._cumulative_oversight_reward = 0.0
        self._done = False
        self._initialized = True
        self._alerts_accurate = False
        self._veto_budget_remaining = self.VETO_BUDGET
        self._flag_budget_remaining = self.FLAG_BUDGET
        self._last_critique = {}

        # Reset investor agent with a fresh persona
        self._investor_agent = InvestorAgent(
            rng=self._rng,
            llm_client=self.investor_llm_client,
            model_name=self.investor_model_name,
        )
        self._investor_agent.reset()
        self._state.investor_state = self._investor_agent.state

        self._baseline_state = run_baseline(task_id, policy_name="naive")

        lessons = self._retrieve_lessons()

        return make_observation(
            self._state,
            self._rng,
            task_description=self._task.description + "\n\n" + describe_generation(self._state.adversary_gen),
            done=False,
            reward=None,
            alerts_accurate=False,
            shared_playbook=[ls.to_obs_dict() for ls in lessons],
            last_oversight_critique={},
        )

    # --- step -------------------------------------------------------------

    def step(
        self,
        action: IncidentAction,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> IncidentObservation:
        if not self._initialized:
            self.reset(task_id="easy_1")
        if self._done:
            self.reset(task_id=self._state.task_id)

        # 0. Treat the incoming action as a Commander proposal
        proposal = CommanderProposal.from_action(action)
        oversight_action: Optional[OversightAction] = kwargs.get("oversight_action")

        # 1. Get Oversight decision
        if oversight_action is None:
            oversight_action = self._oversight_policy(
                proposal, self._state, self._commander_action_history
            )

        # 2. Enforce budgets — downgrade if exhausted
        if oversight_action.decision == OversightDecision.VETO and self._veto_budget_remaining <= 0:
            oversight_action = oversight_action.model_copy(update={"decision": int(OversightDecision.REVISE)})
        if oversight_action.decision == OversightDecision.FLAG_FOR_HUMAN and self._flag_budget_remaining <= 0:
            oversight_action = oversight_action.model_copy(update={"decision": int(OversightDecision.APPROVE)})

        # 3. Build proposal record
        record = ProposalRecord(
            step=self._state.hour,
            proposal=proposal,
            oversight_decision=int(oversight_action.decision),
            oversight_risk_tier=oversight_action.risk_tier,
            oversight_weakness=oversight_action.weakness,
            oversight_counter_action=(
                oversight_action.counter_proposal.action
                if oversight_action.counter_proposal else -1
            ),
            revised=bool(kwargs.get("was_revised", False)),
            final_action=proposal.action,
            final_target=proposal.target_system,
            cited_lessons=list(proposal.cited_lessons),
        )

        # 4. Route on decision
        prev_state = self._state.clone()
        cs = self._state.council_state
        applied = False
        audit_flagged = False
        decision = OversightDecision(oversight_action.decision)

        if decision == OversightDecision.APPROVE:
            cs.approvals += 1
            applied = True
        elif decision == OversightDecision.REVISE:
            cs.revisions += 1
            applied = True
        elif decision == OversightDecision.VETO:
            cs.vetoes += 1
            self._veto_budget_remaining = max(0, self._veto_budget_remaining - 1)
            applied = False
        elif decision == OversightDecision.FLAG_FOR_HUMAN:
            cs.flags += 1
            self._flag_budget_remaining = max(0, self._flag_budget_remaining - 1)
            applied = True
            audit_flagged = True

        # 5. Citations — record attempted citations (even on veto)
        # Track hallucinated ids (cited lesson does not exist in the playbook).
        hallucinated_citations = 0
        if proposal.cited_lessons:
            for lid in proposal.cited_lessons:
                ok = self._playbook.cite(lid)
                if ok:
                    cs.lessons_cited += 1
                else:
                    hallucinated_citations += 1

        # 6. Apply action (if allowed) — bastion vs governance branches
        stamina_cost = 0.0
        team_msgs: List[Dict[str, str]] = []
        governance_result: Dict[str, Any] = {}
        governance_prereq_violations: List[str] = []
        governance_compliance_count = 0

        if applied:
            if is_bastion_action(proposal.action):
                target_name = SYSTEM_NAMES[proposal.target_system]
                governance_prereq_violations = (
                    []
                    if self.disable_governance_prereqs
                    else check_prerequisites(
                        self._state.governance_state,
                        proposal.action,
                        target_name,
                        self._state.hour,
                    )
                )
                if governance_prereq_violations:
                    record_prereq_violations(
                        self._state.governance_state,
                        self._state.hour,
                        governance_prereq_violations,
                        proposal.action,
                        target_name,
                    )
                stamina_cost, self._alerts_accurate = apply_bastion_action(
                    self._state, proposal.action, proposal.target_system, self._rng,
                    method=proposal.method,
                    scope=proposal.scope,
                    rollback_plan=proposal.rollback_plan,
                )
                team_msgs = generate_team_messages(
                    self._state, proposal.action, proposal.target_system, self._rng
                )

            elif is_governance_action(proposal.action):
                target_name = SYSTEM_NAMES[proposal.target_system]
                governance_result = apply_governance_action(
                    self._state.governance_state,
                    proposal.action,
                    target_system=target_name,
                    hour=self._state.hour,
                    severity_arg=proposal.severity_arg,
                    channel_arg=proposal.channel_arg,
                    message_arg=proposal.message_arg,
                    scope_arg=proposal.scope_arg,
                    evidence_arg=proposal.evidence_arg,
                )
                governance_compliance_count = 1
                stamina_cost = 0.02
                self._state.team_stamina = max(0.0, self._state.team_stamina - stamina_cost)
                # If this was a Slack post, see if it answers a pending
                # stakeholder ask (via channel match + non-empty message).
                if proposal.action == int(ActionType.NOTIFY_SLACK_CHANNEL):
                    if not self.disable_stakeholder_events:
                        satisfied = stakeholder_try_respond(
                            self._state.stakeholder_state,
                            channel=proposal.channel_arg,
                            message=proposal.message_arg,
                            hour=self._state.hour,
                        )
                        if satisfied is not None:
                            governance_result["stakeholder_ask_satisfied"] = satisfied.ask_id

                    # Investor channel — handle update regardless of stakeholder flag
                    if proposal.channel_arg in ("investor-relations", "investor_relations"):
                        inv_reply, tier_crossed = self._investor_agent.handle_commander_update(
                            hour=self._state.hour,
                            message_text=proposal.message_arg,
                        )
                        if inv_reply:
                            team_msgs = team_msgs + [inv_reply]
                        if tier_crossed:
                            governance_result["investor_tier_changed"] = self._investor_agent.state.tier()
        else:
            # VETO: small stamina cost for the lost turn
            stamina_cost = 0.02
            self._state.team_stamina = max(0.0, self._state.team_stamina - stamina_cost)

        # 7. Tick any pending process_kill re-compromise events before attacker moves
        recompromised = tick_pending_recompromise(self._state, self._rng)
        for sys_name in recompromised:
            from models import Alert, AlertSeverity
            from dynamics import SYSTEM_IPS
            self._state.alerts.append(Alert(
                source_system=sys_name,
                severity=AlertSeverity.HIGH,
                message=f"Attacker re-established access on {sys_name} — process_kill isolation bypassed via surviving backdoor",
                is_true_positive=True,
                hour=self._state.hour,
                source_ip=SYSTEM_IPS.get(sys_name, "10.0.0.1"),
                dest_ip=SYSTEM_IPS.get(sys_name, "10.0.0.1"),
                mitre_technique="T1543.003",
                mitre_tactic="Persistence",
                process_name="WinSockHelper.exe",
                event_id="EVT-7045",
                confidence=0.88,
            ))

        # 7b. Adversary turn (always advances)
        if self._state.adversary_gen >= 4:
            # Gen 4 — LLM-driven. Use env-configured client + model; fall back
            # to scripted Gen 3 if not available.
            client, model_name = (self.adversary_llm_client, None)
            if client is None:
                client, model_name = make_adversary_client_from_env()
            new_alerts = gen4_adversary_turn(
                self._state, self._rng, client=client, model=model_name
            )
        else:
            new_alerts = adversary_turn(
                self._state, self._rng, generation=self._state.adversary_gen
            )
        self._state.alerts.extend(new_alerts)

        # 7b. Stakeholder pressure events — roll new events and expire overdue
        new_stakeholder_asks = []
        expired_stakeholder_asks = []
        if not self.disable_stakeholder_events:
            new_stakeholder_asks = roll_new_events(
                self._state.stakeholder_state,
                self._rng,
                hour=self._state.hour,
                adversary_gen=self._state.adversary_gen,
                services_disrupted=self._state.services_disrupted,
                data_exfiltrated=self._state.data_exfiltrated,
                management_escalated=self._state.management_escalated,
            )
            expired_stakeholder_asks = expire_overdue_asks(
                self._state.stakeholder_state,
                hour=self._state.hour,
            )
            if new_stakeholder_asks:
                team_msgs = team_msgs + asks_as_team_messages(new_stakeholder_asks)

        # 7c. Investor agent tick — check-ins + anxiety update
        incident_closed = proposal.action == int(ActionType.CLOSE_INCIDENT) and applied
        investor_msgs, investor_pressure = self._investor_agent.tick(
            hour=self._state.hour,
            data_exfiltrated=self._state.data_exfiltrated,
            incident_closed=incident_closed,
        )
        if investor_msgs:
            team_msgs = team_msgs + investor_msgs
        # Sync investor state back to IncidentState for observation / scoring
        self._state.investor_state = self._investor_agent.state

        # 8. Periodic governance checks
        new_gov_violations = periodic_governance_check(
            self._state.governance_state,
            self._state.data_exfiltrated,
            self._state.hour,
        )

        # 9. Management pressure + services update
        if self._state.management_escalated:
            self._state.management_pressure = min(1.0, self._state.management_pressure + 0.05)
        self._state.services_disrupted = sum(
            1 for s in self._state.systems
            if s.name in SERVICE_SYSTEMS and (s.isolated or s.integrity < 0.3)
        )

        # 10. Advance time
        self._state.hour += 1
        self._state.step_count += 1

        # 11. Evaluate council decision
        outcome_correct, critique_precise, counter_succeeded = self._evaluate_council(
            prev_state, self._state, proposal, oversight_action, applied
        )
        record.outcome_correct = outcome_correct

        if decision == OversightDecision.VETO:
            if not outcome_correct:
                cs.correct_vetoes += 1
            else:
                cs.false_vetoes += 1
        if decision == OversightDecision.REVISE and outcome_correct:
            cs.useful_revisions += 1
        if not oversight_action.weakness and decision != OversightDecision.APPROVE:
            cs.vague_critiques += 1
        if oversight_action.counter_proposal and counter_succeeded:
            cs.counter_proposals_adopted += 1
            cs.counter_proposals_succeeded += 1

        # 12. Update trust (skip entirely if trust dynamics disabled → ablation)
        if not self.disable_trust_dynamics:
            self._update_trust(oversight_action, outcome_correct)
            drift_toward_mean(self._state.trust_state)
        self._state.trust_state.snapshot()

        # 13. Write post-mortem lesson (skipped if playbook disabled)
        lesson_utility_delta = 0.0
        if not self.disable_playbook and oversight_action.lesson_text.strip():
            tgt_sys = (
                SYSTEM_NAMES[proposal.target_system]
                if 0 <= proposal.target_system < len(SYSTEM_NAMES) else ""
            )
            alert_conf = -1.0
            for a in reversed(self._state.alerts):
                if a.source_system == tgt_sys and a.confidence > 0:
                    alert_conf = a.confidence
                    break
            auto_tags = make_context_tags(
                adversary_gen=self._state.adversary_gen,
                system_name=tgt_sys,
                alert_confidence=alert_conf,
                extras=oversight_action.lesson_tags,
            )
            self._playbook.write(
                text=oversight_action.lesson_text,
                tags=auto_tags,
                adversary_gen=self._state.adversary_gen,
                task_id=self._state.task_id,
                hour=self._state.hour,
            )

        # 14. Record cited-lesson outcomes
        lesson_was_helpful = False
        if proposal.cited_lessons and outcome_correct:
            lesson_was_helpful = True
            cs.lessons_cited_and_helpful += 1
            for lid in proposal.cited_lessons:
                self._playbook.record_outcome(lid, improved=True)
        elif proposal.cited_lessons and not outcome_correct:
            for lid in proposal.cited_lessons:
                self._playbook.record_outcome(lid, improved=False)

        # 15. Log record
        cs.log(record)

        # 16. Compute rewards
        governance_violation_count = len(governance_prereq_violations) + len(new_gov_violations)

        # Did the Oversight's veto turn out to be correct? (only meaningful on VETO)
        veto_was_correct = None
        if decision == OversightDecision.VETO:
            # outcome_correct is True when the blocked action would have been fine.
            # A "correct" veto is the OPPOSITE — the blocked action would have hurt.
            veto_was_correct = not outcome_correct

        # "Governance chain completed" — destructive action on data system fired
        # WITHOUT any prereq violations this step → the prereq work was done.
        governance_chain_completed = bool(
            applied
            and is_bastion_action(proposal.action)
            and proposal.action in DESTRUCTIVE_BASTION_ACTIONS
            and not governance_prereq_violations
            and SYSTEM_NAMES[proposal.target_system] in DATA_HOLDING_SYSTEMS
        )

        commander_reward = compute_commander_step_reward(
            prev=prev_state,
            curr=self._state,
            stamina_cost=stamina_cost,
            record=record,
            governance_violation_count=governance_violation_count,
            governance_compliance_count=governance_compliance_count,
            governance_chain_completed=governance_chain_completed,
            lesson_was_helpful=lesson_was_helpful,
            veto_was_correct=veto_was_correct,
            hallucinated_citations=hallucinated_citations,
        )
        penalty = compute_penalties(self._state)
        commander_total = commander_reward + penalty
        self._cumulative_commander_reward += commander_total

        oversight_reward = compute_oversight_step_reward(
            record=record,
            outcome_was_correct=outcome_correct,
            critique_was_precise=critique_precise,
            counter_succeeded=counter_succeeded,
            governance_caught=bool(governance_prereq_violations) and not applied,
            lesson_utility_delta=lesson_utility_delta,
            curr_state=self._state,
        )
        self._cumulative_oversight_reward += oversight_reward

        # 17. Stash last critique for next Commander obs (revision UI)
        self._last_critique = {
            "decision": oversight_action.decision,
            "risk_tier": oversight_action.risk_tier,
            "weakness": oversight_action.weakness,
            "missing_evidence": oversight_action.missing_evidence,
            "counter_proposal": (
                oversight_action.counter_proposal.model_dump()
                if oversight_action.counter_proposal else None
            ),
        }

        # 18. Track Commander action history
        self._commander_action_history.append({
            "hour": self._state.hour,
            "action": ACTION_NAMES.get(proposal.action, str(proposal.action)),
            "target": proposal.target_system,
            "justification": proposal.justification[:120],
            "decision": decision.name,
            "outcome_correct": outcome_correct,
        })

        # 19. Termination check
        done = False

        # Per-step system snapshot for dashboard replay
        systems_snapshot = {
            s.name: {
                "compromised": s.compromised,
                "isolated": s.isolated,
                "investigated": s.investigated,
                "has_backdoor": s.has_backdoor,
                "integrity": round(s.integrity, 2),
                "criticality": round(getattr(s, "criticality", 0.5), 2),
            }
            for s in self._state.systems
        }

        # Alerts fired this step (new alerts from adversary turn)
        step_alerts = [
            {
                "severity": a.severity.name if hasattr(a.severity, "name") else str(a.severity),
                "system": a.source_system,
                "message": a.message,
                "mitre": getattr(a, "mitre_technique", ""),
                "mitre_tactic": getattr(a, "mitre_tactic", ""),
                "event_id": getattr(a, "event_id", ""),
                "confidence": round(getattr(a, "confidence", 0.0), 2),
                "source_ip": getattr(a, "source_ip", ""),
                "dest_ip": getattr(a, "dest_ip", ""),
                "process": getattr(a, "process_name", ""),
                "is_true_positive": getattr(a, "is_true_positive", True),
                "hour": getattr(a, "hour", self._state.hour),
            }
            for a in new_alerts
        ]

        # Investor messages fired this step
        investor_step_messages = [
            {
                "hour": m.hour,
                "direction": m.direction,
                "text": m.text,
                "anxiety_before": round(m.anxiety_before, 3),
                "anxiety_after": round(m.anxiety_after, 3),
            }
            for m in self._investor_agent.state.messages
            if m.hour == self._state.hour
        ]

        # Playbook top lessons snapshot (by utility, no tag filter)
        top_lessons = [
            {"lesson_id": ls.lesson_id, "text": ls.text[:120], "utility": round(ls.utility, 3),
             "citations": ls.citations, "wins": ls.wins, "losses": ls.losses}
            for ls in sorted(self._playbook.all(), key=lambda l: l.utility, reverse=True)[:8]
        ]

        info: Dict[str, Any] = {
            "hour": self._state.hour,
            "action_name": ACTION_NAMES.get(proposal.action, str(proposal.action)),
            "stamina_cost": round(stamina_cost, 3),
            "oversight_decision": decision.name,
            "oversight_risk_tier": oversight_action.risk_tier,
            "applied": applied,
            "audit_flagged": audit_flagged,
            "step_reward": round(commander_total, 4),
            "oversight_reward": round(oversight_reward, 4),
            "penalty": round(penalty, 4),
            "governance_prereq_violations": governance_prereq_violations,
            "governance_new_violations": new_gov_violations,
            "governance_result": governance_result,
            "stakeholder_new_asks": [a.model_dump() for a in new_stakeholder_asks],
            "stakeholder_expired_asks": [a.ask_id for a in expired_stakeholder_asks],
            "stakeholder_reputation": round(self._state.stakeholder_state.reputation, 3),
            "outcome_correct": outcome_correct,
            "critique_precise": critique_precise,
            # Dashboard replay fields
            "systems_snapshot": systems_snapshot,
            "step_alerts": step_alerts,
            "investor_step_messages": investor_step_messages,
            "investor_anxiety": round(self._investor_agent.state.anxiety, 3),
            "investor_tier": self._investor_agent.state.tier(),
            "investor_persona": PERSONAS[self._investor_agent.state.persona_id]["name"]
                if hasattr(self._investor_agent.state, "persona_id") else "",
            "data_exfiltrated": round(self._state.data_exfiltrated, 4),
            "team_stamina": round(self._state.team_stamina, 3),
            "management_pressure": round(self._state.management_pressure, 3),
            "services_disrupted": self._state.services_disrupted,
            "trust_snapshot": self._state.trust_state.as_dict(),
            "playbook_snapshot": top_lessons,
        }

        termination_reason = ""
        if self._state.data_exfiltrated >= 1.0:
            done = True
            termination_reason = "total_data_breach"
            info["termination_reason"] = termination_reason
        elif self._state.hour >= self._task.max_hours:
            done = True
            termination_reason = "time_expired"
            info["termination_reason"] = termination_reason

        if done:
            self._done = True
            self._playbook.save()
            had_p1 = any(t.severity == 1 for t in self._state.governance_state.tickets)
            investor_score = self._investor_agent.investor_score()
            final_scores = compute_final_score(
                state=self._state,
                weights=self._task.scoring_weights,
                proposals_made=self._state.council_state.total_proposals,
                max_proposals=self._task.max_hours,
                per_gen_scores=None,
                had_p1=had_p1,
                termination_reason=termination_reason,
                investor_score=investor_score,
            )
            baseline_score = compute_bastion_v1_final_score(
                self._baseline_state, self._task.scoring_weights
            )
            comparison = compute_baseline_comparison(
                self._state, self._baseline_state, self._task.scoring_weights
            )
            info["final_scores"] = final_scores
            info["baseline_final_score"] = round(baseline_score, 4)
            info["comparison_score"] = round(comparison, 4)
            info["cumulative_commander_reward"] = round(self._cumulative_commander_reward, 4)
            info["cumulative_oversight_reward"] = round(self._cumulative_oversight_reward, 4)
            info["data_exfiltrated"] = round(self._state.data_exfiltrated, 4)
            info["attacker_progress"] = round(self._state.attacker_progress, 4)
            info["adversary_gen"] = self._state.adversary_gen
            info["council_summary"] = {
                "total_proposals": cs.total_proposals,
                "approvals": cs.approvals,
                "revisions": cs.revisions,
                "vetoes": cs.vetoes,
                "flags": cs.flags,
                "correct_vetoes": cs.correct_vetoes,
                "false_vetoes": cs.false_vetoes,
                "useful_revisions": cs.useful_revisions,
                "vague_critiques": cs.vague_critiques,
                "counter_proposals_adopted": cs.counter_proposals_adopted,
                "counter_proposals_succeeded": cs.counter_proposals_succeeded,
                "lessons_cited": cs.lessons_cited,
                "lessons_cited_and_helpful": cs.lessons_cited_and_helpful,
            }
            info["trust_final"] = self._state.trust_state.as_dict()
            info["governance_final"] = self._state.governance_state.snapshot()
            info["stakeholder_final"] = self._state.stakeholder_state.snapshot()
            info["investor_final"] = self._investor_agent.state.snapshot()
            info["investor_score"] = round(investor_score, 4)
            info["forensic_report"] = generate_forensic_report(self._state)

        # 20. Build Commander observation for next step
        lessons = self._retrieve_lessons()
        obs = make_observation(
            self._state,
            self._rng,
            task_description=self._task.description if not done else "",
            done=done,
            reward=commander_total,
            alerts_accurate=self._alerts_accurate,
            team_messages=team_msgs,
            shared_playbook=[ls.to_obs_dict() for ls in lessons],
            last_oversight_critique=self._last_critique,
        )
        obs.oversight_reward = oversight_reward
        obs.metadata = info
        return obs

    # --- helpers ----------------------------------------------------------

    def _retrieve_lessons(self, max_items: int = 6):
        if self.disable_playbook:
            return []
        tags = make_context_tags(
            adversary_gen=self._state.adversary_gen,
            system_name="",
            alert_confidence=-1.0,
        )
        return self._playbook.retrieve(tags=tags, max_items=max_items)

    def _evaluate_council(
        self,
        prev_state: IncidentState,
        curr_state: IncidentState,
        proposal: CommanderProposal,
        oversight_action: OversightAction,
        applied: bool,
    ) -> tuple:
        """
        Proxy evaluation of the council decision.
          - outcome_correct: did this step net-improve defender metrics?
          - critique_precise: did Oversight's named weakness match a known
            anti-pattern (uninvestigated target, missing prerequisite)?
          - counter_succeeded: did the counter-proposal suggest a sensible
            next action (e.g., investigate before isolate)?
        """
        delta_exfil = curr_state.data_exfiltrated - prev_state.data_exfiltrated
        delta_compromised = curr_state.compromised_count - prev_state.compromised_count
        net_positive = (delta_exfil <= 0.02) and (delta_compromised <= 0)

        if applied:
            outcome_correct = net_positive
        else:
            checks = compute_policy_checks(prev_state, proposal)
            obviously_bad = (
                checks["action_targets_uninvestigated_system"]
                and checks["action_irreversible"]
            ) or bool(checks["prerequisites_missing"])
            outcome_correct = not obviously_bad

        weakness_given = bool(oversight_action.weakness.strip())
        known_anti_pattern = (
            proposal.action in DESTRUCTIVE_BASTION_ACTIONS
            and compute_policy_checks(prev_state, proposal)["action_targets_uninvestigated_system"]
        )
        critique_precise = weakness_given and (
            known_anti_pattern or len(oversight_action.missing_evidence) > 0
        )

        counter_succeeded = False
        if oversight_action.counter_proposal is not None:
            cp = oversight_action.counter_proposal
            if cp.action == int(ActionType.INVESTIGATE_SYSTEM):
                try:
                    s = prev_state.get_system_by_idx(cp.target_system)
                    if not s.investigated and s.name in DATA_HOLDING_SYSTEMS:
                        counter_succeeded = True
                except Exception:
                    pass

        return outcome_correct, critique_precise, counter_succeeded

    def _update_trust(self, oversight_action: OversightAction, outcome_correct: bool) -> None:
        ts = self._state.trust_state
        decision = OversightDecision(oversight_action.decision)

        # Oversight's trust in Commander
        if outcome_correct:
            update_trust_o2c(ts, "correct")
        else:
            update_trust_o2c(ts, "rework_needed")

        # Commander's trust in Oversight
        if decision == OversightDecision.VETO:
            update_trust_c2o(ts, "veto_correct" if not outcome_correct else "veto_wrong")
        elif decision == OversightDecision.APPROVE:
            update_trust_c2o(ts, "approve_correct" if outcome_correct else "approve_wrong")
        elif decision == OversightDecision.REVISE:
            update_trust_c2o(ts, "demand_useful" if outcome_correct else "demand_nitpick")

    @property
    def state(self) -> IncidentState:
        return self._state