inference.py

"""
Citadel — Inference Script
===========================
MANDATORY env vars (hackathon requirements):
    API_BASE_URL   The API endpoint for the LLM.
    MODEL_NAME     The model identifier to use for inference.
    HF_TOKEN       Your Hugging Face / API key.

STDOUT FORMAT (strict): [START], [STEP], [END]

Runs two LLMs in a council protocol over three tasks (easy_1, medium_1,
hard_1). Each step:
  1. Commander LLM proposes (action + target + justification + cited lessons).
  2. Oversight LLM critiques (decision + risk_tier + weakness + counter + lesson).
  3. On REVISE, Commander revises ONCE with the critique in context.
  4. Env applies the (possibly revised) proposal through the approved route.

Fallback chain: Docker image → HF Space → LocalEnv.
"""

from __future__ import annotations

import asyncio
import json
import os
import re
import textwrap
import time
from pathlib import Path
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

from openai import OpenAI

from models import (
    IncidentAction,
    IncidentObservation,
    OversightAction,
    OversightDecision,
    CommanderProposal,
    CounterProposal,
    ACTION_NAMES,
    NUM_ACTIONS,
    SYSTEM_NAMES,
)
from environment import CitadelEnvironment
from recorder import RunRecorder, make_run_root, write_run_index


# ---------------------------------------------------------------------------
# Configuration (mandatory env vars per hackathon spec)
# ---------------------------------------------------------------------------

API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN")

LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
API_KEY = HF_TOKEN

BENCHMARK = "citadel"
MAX_STEPS = 12
TEMPERATURE = 0.3
MAX_TOKENS_COMMANDER = 500
MAX_TOKENS_OVERSIGHT = 450
TASKS = ["easy_1", "medium_1", "hard_1"]


# ---------------------------------------------------------------------------
# Local environment wrapper (matches the EnvClient StepResult interface)
# ---------------------------------------------------------------------------

@dataclass
class StepResult:
    observation: IncidentObservation
    reward: Optional[float] = None
    done: bool = False


class LocalEnv:
    """Wraps CitadelEnvironment to match the async EnvClient interface."""

    def __init__(self, investor_llm_client: Optional[Any] = None, investor_model_name: str = "") -> None:
        self._env = CitadelEnvironment(
            investor_llm_client=investor_llm_client,
            investor_model_name=investor_model_name,
        )

    async def reset(self, **kwargs: Any) -> StepResult:
        obs = self._env.reset(**kwargs)
        return StepResult(observation=obs, reward=obs.reward, done=obs.done)

    async def step(
        self,
        action: IncidentAction,
        oversight_action: Optional[OversightAction] = None,
        was_revised: bool = False,
    ) -> StepResult:
        obs = self._env.step(action, oversight_action=oversight_action, was_revised=was_revised)
        return StepResult(observation=obs, reward=obs.reward, done=obs.done)

    async def close(self) -> None:
        pass


# ---------------------------------------------------------------------------
# Structured stdout logging (hackathon-required format)
# ---------------------------------------------------------------------------

def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}", flush=True)


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)


# ---------------------------------------------------------------------------
# System prompts — Commander + Oversight
# ---------------------------------------------------------------------------

COMMANDER_SYSTEM_PROMPT = textwrap.dedent("""\
    You are the AI INCIDENT COMMANDER in a Security Operations Center during a live cyberattack.
    You work in a two-agent council with an AI OVERSIGHT reviewer. Your goal: contain the incident
    while following enterprise governance.

    ## Network (8 systems, indexed 0-7):
    0: web_server    1: app_server    2: database      3: file_server
    4: email_server  5: workstations  6: backup_server  7: firewall

    ## Action space — 18 actions
    Incident response (0-9):
      0: investigate_system      1: isolate_system       2: patch_vulnerability
      3: restore_from_backup     4: analyze_alerts       5: deploy_monitoring
      6: escalate_to_management  7: block_external_traffic  8: hunt_threat
      9: coordinate_team
    Governance (10-17):
     10: open_servicenow_incident (severity_arg=P1..P4)
     11: request_cab_approval     (severity_arg=action, target_system=system)
     12: notify_slack_channel     (channel_arg, message_arg)
     13: log_to_sox_audit         (severity_arg=action, target_system=system, evidence_arg)
     14: page_oncall              (channel_arg=team, severity_arg=severity)
     15: notify_data_owner        (target_system=system)
     16: start_legal_hold         (scope_arg)
     17: close_incident           (evidence_arg=summary)

    ## Governance pre-requisites (compliance is scored!)
    - BEFORE isolate() on a data system: request_cab_approval AND notify_data_owner first.
    - Before block_external_traffic: post to sec-leadership on Slack first.
    - P1 incident needs ServiceNow + PagerDuty + sec-leadership Slack within 2 hours.
    - If data_exfiltrated > 20%, a 72-hour GDPR clock starts — start_legal_hold in time.

    ## Investor Relations (#investor-relations Slack channel)
    Nervous investors are watching. You MUST keep them calm and informed.
    - Use notify_slack_channel(channel="investor-relations", message="...") to update them.
    - Keep updates calm, specific, and confident. Avoid words like "breach", "critical", "unknown".
    - Vague updates ("working on it") increase investor anxiety and trigger board escalation.
    - Good update: "Incident contained to web_server tier. No customer data accessed. ETA 2h to resolve."
    - Bad update:  "We have a major breach and are still investigating."
    - Investor anxiety appears in your observation — high anxiety = board pressure events incoming.
    - Post at least once every 3 hours to prevent silence penalty.

    ## Shared Playbook
    Both you and Oversight see a PLAYBOOK of lessons from past episodes (id like L-0017).
    If a lesson applies to your current situation, include its id in `cited_lessons`.

    ## Council loop
    Oversight returns a structured critique. On REVISE you get ONE chance to revise.

    ## Richer action payloads — method, scope, rollback_plan
    For isolate_system, choose method:
      "firewall_acl"   — reversible, service stays partially up, backdoor survives (use when unsure)
      "network_unplug" — complete cut, fastest containment, full service loss, backdoor severed
      "process_kill"   — kills attacker processes only, service stays up, 40% chance attacker returns in 1-2h

    For patch_vulnerability, choose method:
      "hotpatch"       — no restart, 60% clean chance, minimal disruption
      "cold_patch"     — full restart, 90% clean chance, brief service hit
      "virtual_patch"  — WAF rule blocks exploit, no binary change, only for web/app/database/email

    For deploy_monitoring, choose method:
      "process_events"  — catches lateral movement (default)
      "network_traffic" — catches exfil/C2 beacons, also improves neighbors
      "auth_events"     — catches credential abuse, low noise
      "full_endpoint"   — catches everything, -5% integrity/hour on target

    scope: IP range or process name to constrain the action (e.g. "10.1.3.30/32", "lsass.exe")
    rollback_plan: how to undo if wrong — REQUIRED for isolate/block/patch (improves effectiveness +10%)

    ## Output — ONLY a compact JSON object:
    {"action": <0-17>, "target": <0-7>, "justification": "<reason>", "cited_lessons": ["L-id", ...],
     "method": "<method>", "scope": "<scope>", "rollback_plan": "<how to undo>",
     "severity_arg": <1-4>, "channel_arg": "...", "message_arg": "...",
     "scope_arg": "...", "evidence_arg": "..."}
""")


OVERSIGHT_SYSTEM_PROMPT = textwrap.dedent("""\
    You are the AI OVERSIGHT reviewer in a Security Operations Center. The Commander proposes
    actions during a live cyberattack. Your job is NOT to rubber-stamp and NOT to micromanage:
    CRITIQUE WITH SUBSTANCE — name the weakest assumption, ask for missing evidence, and when
    the proposal is clearly wrong, offer a concrete counter-action.

    ## Decisions
    0: APPROVE         — proposal is sound
    1: REVISE          — fixable weakness; Commander revises once (you will see the revision)
    2: VETO            — net harm; block it (limited veto budget — use sparingly)
    3: FLAG_FOR_HUMAN  — ambiguous / high-risk; execute + audit (limited budget)

    ## Your observation structure (each step)
    1. INDEPENDENT ALERT DIGEST — the raw SIEM feed before seeing Commander's proposal.
       Form your OWN threat picture here first. Note which systems look compromised,
       which alerts look like false flags, what the attacker appears to be doing.
    2. COMMANDER'S PROPOSAL — what Commander wants to do, with method/scope/rollback.
    3. EPISODE HISTORY — your own past decisions and what actually happened after each.
       Use this to reason about the incident arc, not just the current step in isolation.
    4. POLICY CHECKS — pre-computed governance prerequisite violations.

    ## Anti-patterns to watch for
    - Destructive action (isolate, block_external_traffic) on an UNINVESTIGATED system.
    - Destructive action missing governance prerequisites (no CAB, no data_owner notify).
    - Justification that doesn't cite any alert / log / lesson.
    - Isolating clean systems on false-positive alerts (Gen 2 false flags).
    - Gen 3 deceptive alerts: low-severity "all clear" on a system with prior HIGH alerts — don't be fooled.
    - isolate with method=process_kill on a backdoored system — attacker will return in 1-2h.
    - patch with method=virtual_patch on a non-network-facing system (workstations, backup) — no-op.
    - deploy_monitoring with method=full_endpoint when system integrity is already low — accelerates degradation.
    - No rollback_plan on isolate/patch/block — signals Commander didn't think about failure modes.
    - Repeating the same action on the same system that already failed — check episode history.

    ## Revision enforcement
    When you issue REVISE with a counter_proposal, Commander must address it.
    If Commander's revision ignores your counter_proposal entirely (different action AND different target),
    treat it as if they submitted a new unreviewed proposal — escalate to VETO automatically.

    ## Post-mortem lesson
    Write a one-line generalizable rule (lesson_text). Lessons persist across episodes.
    Good: "Alerts below 30% confidence on data systems require investigate() before isolate()."
    Bad:  "Be careful." (too vague)

    ## Output — ONLY a compact JSON object:
    {"decision": <0-3>, "risk_tier": <1-5>, "weakness": "...",
     "missing_evidence": ["..."], "counter_proposal": {"action":<0-17>,"target":<0-7>,"rationale":"..."} | null,
     "lesson_text": "<one-line rule>", "lesson_tags": ["..."]}
""")


# ---------------------------------------------------------------------------
# Observation formatting
# ---------------------------------------------------------------------------

def format_commander_observation(obs: dict, step: int, history: List[str], revision_context: str = "") -> str:
    parts: List[str] = []

    if step == 0:
        desc = obs.get("task_description", "")
        if desc:
            parts.append(f"## Incident Brief\n{desc}\n")

    parts.append(f"## Hour {obs.get('hour', 0)} (remaining: {obs.get('hours_remaining', 12)})")
    parts.append(f"- Breach: {obs.get('estimated_breach_severity', 'unknown')} | Data at risk: {obs.get('estimated_data_at_risk', 0):.0%}")
    parts.append(f"- Services disrupted: {obs.get('services_disrupted', 0)}/{obs.get('services_total', 4)} | Stamina: {obs.get('team_stamina', 1.0):.0%}")
    parts.append(f"- External blocked: {obs.get('external_blocked', False)} | Mgmt escalated: {obs.get('management_escalated', False)}")

    trust = obs.get("trust_summary", {})
    if trust:
        parts.append(
            f"- Trust: self→O={trust.get('trust_commander_in_oversight', 0):.2f}, "
            f"O→self={trust.get('trust_oversight_in_commander', 0):.2f}"
        )

    gov = obs.get("governance_summary", {})
    if gov:
        parts.append("\n## Governance")
        parts.append(f"  tickets_open={len(gov.get('open_tickets', []))} cab={gov.get('cab_approvals_count', 0)} sox={gov.get('sox_log_count', 0)} slack={gov.get('slack_posts_count', 0)} pages={gov.get('pages_count', 0)}")
        owners = gov.get("data_owners_notified", [])
        if owners:
            parts.append(f"  data_owners_notified: {owners}")
        if gov.get("gdpr_clock_hours_remaining") is not None:
            parts.append(f"  gdpr_clock={gov['gdpr_clock_hours_remaining']}h legal_hold={gov.get('legal_hold_active')}")
        if gov.get("violations_count", 0) > 0:
            parts.append(f"  VIOLATIONS: {gov['violations_count']}")

    systems = obs.get("systems_visible", [])
    if systems:
        parts.append("\n## Systems")
        for s in systems:
            bits = []
            comp = s.get("compromised", "unknown")
            bits.append(f"comp={'Y' if comp is True else 'N' if comp is False else '?'}")
            if s.get("isolated"):
                bits.append("ISO")
            if s.get("investigated"):
                bits.append("inv")
            if s.get("patched"):
                bits.append("pat")
            bits.append(f"int={s.get('integrity', 1.0):.0%}")
            idx = SYSTEM_NAMES.index(s["name"]) if s["name"] in SYSTEM_NAMES else 0
            parts.append(f"  [{idx}] {s['name']:14s} | {' '.join(bits)}")

    alerts = obs.get("alert_queue", [])
    if alerts:
        parts.append("\n## SIEM Alerts (recent 4)")
        for a in alerts[-4:]:
            eid = a.get("event_id", "")
            conf = f" c={a['confidence']:.0%}" if a.get("confidence") else ""
            parts.append(f"  [{a.get('severity', '?'):8s}] {eid} {a.get('source_system', '?')}: {a.get('message', '')}{conf}")

    lessons = obs.get("shared_playbook", [])
    if lessons:
        parts.append("\n## Shared Playbook")
        for ls in lessons[:6]:
            parts.append(f"  {ls['id']} (u={ls['utility']:+.2f}): {ls['text']}")

    investor = obs.get("investor_summary", {})
    if investor:
        tier = investor.get("tier", "?")
        anxiety = investor.get("anxiety", 0)
        persona = investor.get("persona", "Investor")
        last_upd = investor.get("last_update_hour", -1)
        parts.append(f"\n## #investor-relations — {persona}")
        parts.append(f"  anxiety={anxiety:.2f} tier={tier} last_update=hour{last_upd}")
        if tier in ("ALARMED", "PANIC"):
            parts.append(f"  ⚠ POST UPDATE NOW — investor is {tier}, board escalation imminent")
        elif tier == "CONCERNED":
            parts.append(f"  Post a specific update to calm them (channel='investor-relations')")

    team = obs.get("team_messages", [])
    if team:
        parts.append("\n## Team Comms")
        for m in team[-2:]:
            parts.append(f"  [{m.get('from', '?')}] {m.get('message', '')[:180]}")

    if history:
        parts.append("\n## Recent actions")
        for h in history[-3:]:
            parts.append(f"  {h}")

    if revision_context:
        parts.append(f"\n## Oversight critique to address\n{revision_context}")

    parts.append('\nRespond ONLY with JSON: {"action":<0-17>,"target":<0-7>,"justification":"...","cited_lessons":[...]}')
    return "\n".join(parts)


def format_oversight_observation(oobs: dict) -> str:
    parts: List[str] = []

    # Section 1 — Independent alert digest (raw SIEM before seeing proposal)
    raw_alerts = oobs.get("raw_alert_digest", [])
    if raw_alerts:
        parts.append("## 1. INDEPENDENT ALERT DIGEST (form your threat picture BEFORE reading the proposal)")
        for a in raw_alerts[-6:]:
            eid = a.get("event_id", "")
            conf = f" c={a['confidence']:.0%}" if a.get("confidence") else ""
            sev = a.get("severity", "?")
            # severity may be int (AlertSeverity enum) or string
            sev_str = ["low", "medium", "high", "critical"][sev] if isinstance(sev, int) and 0 <= sev <= 3 else str(sev)
            parts.append(
                f"  [{sev_str:8s}] {eid} {a.get('source_system', '?')}: "
                f"{a.get('message', '')}{conf}"
            )
    else:
        parts.append("## 1. INDEPENDENT ALERT DIGEST\n  (no alerts in queue)")

    # Section 2 — Commander's proposal
    pa = oobs.get("proposed_action", {})
    parts.append("\n## 2. COMMANDER'S PROPOSAL")
    parts.append(f"  {pa.get('action_name', '?')}({pa.get('action', '?')}) on {pa.get('target_name', '?')}[{pa.get('target_system', '?')}]")
    if pa.get("method"):
        parts.append(f"  method={pa['method']}  scope={pa.get('scope', '') or '(all)'}  rollback={pa.get('rollback_plan', '') or '(none specified)'}")
    if pa.get("severity_arg", 2) != 2 or pa.get("message_arg"):
        parts.append(f"  sev={pa.get('severity_arg', 2)} ch={pa.get('channel_arg')} msg={(pa.get('message_arg') or '')[:100]}")
    parts.append(f"  justification: {(oobs.get('justification') or '')[:400]}")
    cited = oobs.get("cited_lessons", [])
    if cited:
        parts.append(f"  cited_lessons: {cited}")

    # Section 3 — Episode history (Oversight's own past decisions)
    ep_hist = oobs.get("oversight_episode_history", [])
    if ep_hist:
        parts.append("\n## 3. EPISODE HISTORY (your past decisions this incident)")
        for entry in ep_hist[-6:]:
            outcome = entry.get("outcome", "?")
            parts.append(
                f"  Hour {entry.get('hour', '?')}: {entry.get('decision', '?')} | "
                f"{entry.get('action_name', '?')}({entry.get('target', '?')}) → {outcome}"
            )
    else:
        parts.append("\n## 3. EPISODE HISTORY\n  (first step — no history yet)")

    # Section 4 — Policy checks
    pc = oobs.get("policy_checks", {})
    parts.append("\n## 4. POLICY CHECKS")
    if pc:
        for k, v in pc.items():
            parts.append(f"  {k}: {v}")
    else:
        parts.append("  (no violations detected)")

    parts.append(f"\n## Budgets: veto={oobs.get('veto_budget_remaining', '?')} flag={oobs.get('flag_budget_remaining', '?')}")

    trust = oobs.get("trust_summary", {})
    if trust:
        parts.append(f"## Trust: self→C={trust.get('trust_oversight_in_commander', 0):.2f} C→self={trust.get('trust_commander_in_oversight', 0):.2f}")

    lessons = oobs.get("shared_playbook", [])
    if lessons:
        parts.append("\n## Playbook (top lessons)")
        for ls in lessons[:5]:
            parts.append(f"  {ls['id']} (u={ls['utility']:+.2f}): {ls['text']}")

    parts.append('\nRespond ONLY with JSON: {"decision":<0-3>,"risk_tier":<1-5>,"weakness":"...","missing_evidence":[...],"counter_proposal":{...}|null,"lesson_text":"...","lesson_tags":[...]}')
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# Response parsing
# ---------------------------------------------------------------------------

def _extract_json_block(text: str) -> Optional[dict]:
    m = re.search(r"\{[\s\S]*\}", text)
    if not m:
        return None
    raw = m.group()
    try:
        return json.loads(raw)
    except Exception:
        pass
    depth = 0
    for i, ch in enumerate(raw):
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                try:
                    return json.loads(raw[: i + 1])
                except Exception:
                    return None
    return None


def parse_commander_response(text: str) -> IncidentAction:
    data = _extract_json_block(text) or {}
    try:
        action = int(data.get("action", 9))
    except Exception:
        action = 9
    try:
        target = int(data.get("target", 0))
    except Exception:
        target = 0
    action = max(0, min(NUM_ACTIONS - 1, action))
    target = max(0, min(len(SYSTEM_NAMES) - 1, target))

    sev_raw = data.get("severity_arg", 2)
    try:
        severity = max(1, min(4, int(sev_raw)))
    except Exception:
        severity = 2

    return IncidentAction(
        action=action,
        target_system=target,
        justification=str(data.get("justification", ""))[:1000],
        cited_lessons=[str(x) for x in (data.get("cited_lessons") or []) if x][:6],
        method=str(data.get("method", ""))[:32],
        scope=str(data.get("scope", ""))[:200],
        rollback_plan=str(data.get("rollback_plan", ""))[:300],
        severity_arg=severity,
        channel_arg=str(data.get("channel_arg", "sec-ops"))[:64],
        message_arg=str(data.get("message_arg", ""))[:400],
        scope_arg=str(data.get("scope_arg", ""))[:200],
        evidence_arg=str(data.get("evidence_arg", ""))[:400],
    )


def parse_oversight_response(text: str) -> OversightAction:
    data = _extract_json_block(text) or {}
    try:
        decision = int(data.get("decision", 0))
    except Exception:
        decision = 0
    decision = max(0, min(3, decision))
    try:
        risk = int(data.get("risk_tier", 2))
    except Exception:
        risk = 2
    risk = max(1, min(5, risk))

    cp_data = data.get("counter_proposal")
    cp: Optional[CounterProposal] = None
    if isinstance(cp_data, dict):
        try:
            cp = CounterProposal(
                action=max(0, min(NUM_ACTIONS - 1, int(cp_data.get("action", 0)))),
                target_system=max(0, min(len(SYSTEM_NAMES) - 1, int(cp_data.get("target", 0)))),
                rationale=str(cp_data.get("rationale", ""))[:400],
            )
        except Exception:
            cp = None

    return OversightAction(
        decision=decision,
        risk_tier=risk,
        weakness=str(data.get("weakness", ""))[:400],
        missing_evidence=[str(x)[:160] for x in (data.get("missing_evidence") or []) if x][:5],
        counter_proposal=cp,
        lesson_text=str(data.get("lesson_text", ""))[:240],
        lesson_tags=[str(x)[:40] for x in (data.get("lesson_tags") or []) if x][:6],
    )


# ---------------------------------------------------------------------------
# Helper — flatten governance events from step metadata for recorder
# ---------------------------------------------------------------------------

def _extract_governance_events(meta: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Turn governance_result dict + violations into a flat list of events for the dashboard."""
    events: List[Dict[str, Any]] = []
    result = meta.get("governance_result") or {}
    hour = meta.get("hour", 0)

    # Compliance hits from governance_result keys
    for key, val in result.items():
        events.append({"kind": key, "detail": val, "hour": hour, "type": "compliance"})

    # Prerequisite violations
    for v in meta.get("governance_prereq_violations") or []:
        events.append({"kind": v, "hour": hour, "type": "violation"})

    # Periodic governance violations
    for v in meta.get("governance_new_violations") or []:
        events.append({"kind": v, "hour": hour, "type": "violation"})

    return events


# ---------------------------------------------------------------------------
# Council loop — one step (proposal → critique → optional revise → apply)
# ---------------------------------------------------------------------------

async def council_step(
    client: OpenAI,
    env,
    commander_obs: dict,
    commander_history: List[str],
    step_idx: int,
    recorder: Optional[RunRecorder] = None,
    oversight_history: Optional[List[dict]] = None,
) -> Tuple[StepResult, IncidentAction, OversightAction, bool]:
    """One full council step. Returns (step_result, final_action, oversight_action, was_revised)."""

    # --- 1. Commander proposes -----------------------------------------------
    user_msg = format_commander_observation(commander_obs, step_idx, commander_history)
    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": COMMANDER_SYSTEM_PROMPT},
                {"role": "user", "content": user_msg},
            ],
            max_tokens=MAX_TOKENS_COMMANDER,
            temperature=TEMPERATURE,
            stream=False,
        )
        cmd_text = (completion.choices[0].message.content or "").strip()
    except Exception as exc:
        print(f"[DEBUG] Commander LLM error: {exc}", flush=True)
        cmd_text = '{"action": 9, "target": 0, "justification": "API error fallback"}'

    proposal_action = parse_commander_response(cmd_text)
    _commander_prompt_for_recorder = user_msg
    _commander_raw_for_recorder = cmd_text

    # --- 2. Oversight critiques ----------------------------------------------
    # Option B: build independent alert digest from raw SIEM feed BEFORE seeing proposal
    raw_alerts = commander_obs.get("alert_queue", [])

    oobs = {
        "proposed_action": {
            "action": proposal_action.action,
            "action_name": ACTION_NAMES.get(proposal_action.action, str(proposal_action.action)),
            "target_system": proposal_action.target_system,
            "target_name": SYSTEM_NAMES[proposal_action.target_system] if 0 <= proposal_action.target_system < len(SYSTEM_NAMES) else "",
            "method": proposal_action.method,
            "scope": proposal_action.scope,
            "rollback_plan": proposal_action.rollback_plan,
            "severity_arg": proposal_action.severity_arg,
            "channel_arg": proposal_action.channel_arg,
            "message_arg": proposal_action.message_arg,
            "scope_arg": proposal_action.scope_arg,
            "evidence_arg": proposal_action.evidence_arg,
        },
        "justification": proposal_action.justification,
        "cited_lessons": list(proposal_action.cited_lessons),
        "commander_observation": commander_obs,
        "commander_action_history": [{"action": h} for h in commander_history[-4:]],
        "policy_checks": {},
        "veto_budget_remaining": 2,
        "flag_budget_remaining": 2,
        "shared_playbook": commander_obs.get("shared_playbook", []),
        "trust_summary": commander_obs.get("trust_summary", {}),
        "governance_summary": commander_obs.get("governance_summary", {}),
        # Option A: episode history for Oversight
        "oversight_episode_history": oversight_history or [],
        # Option B: raw alert feed for independent digest
        "raw_alert_digest": raw_alerts,
    }
    ov_user_msg = format_oversight_observation(oobs)
    try:
        completion_o = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": OVERSIGHT_SYSTEM_PROMPT},
                {"role": "user", "content": ov_user_msg},
            ],
            max_tokens=MAX_TOKENS_OVERSIGHT,
            temperature=TEMPERATURE,
            stream=False,
        )
        oversight_text = (completion_o.choices[0].message.content or "").strip()
    except Exception as exc:
        print(f"[DEBUG] Oversight LLM error: {exc}", flush=True)
        oversight_text = '{"decision": 0, "risk_tier": 2}'

    oversight_action = parse_oversight_response(oversight_text)
    _oversight_prompt_for_recorder = ov_user_msg
    _oversight_raw_for_recorder = oversight_text

    # --- 3. Revision cycle (one chance) --------------------------------------
    was_revised = False
    _revision_prompt = None
    _revision_raw = None
    if oversight_action.decision == int(OversightDecision.REVISE):
        revision_ctx = (
            f"DECISION: REVISE\n"
            f"risk_tier: {oversight_action.risk_tier}\n"
            f"weakness: {oversight_action.weakness}\n"
            f"missing_evidence: {oversight_action.missing_evidence}"
        )
        if oversight_action.counter_proposal:
            revision_ctx += (
                f"\ncounter_proposal: action={oversight_action.counter_proposal.action}, "
                f"target={oversight_action.counter_proposal.target_system}, "
                f"rationale={oversight_action.counter_proposal.rationale}"
            )
        revision_ctx += "\n\nAddress the weakness above — strengthen justification and adjust action if needed."

        rev_msg = format_commander_observation(
            commander_obs, step_idx, commander_history, revision_context=revision_ctx
        )
        try:
            completion_r = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": COMMANDER_SYSTEM_PROMPT},
                    {"role": "user", "content": rev_msg},
                ],
                max_tokens=MAX_TOKENS_COMMANDER,
                temperature=TEMPERATURE,
                stream=False,
            )
            rev_text = (completion_r.choices[0].message.content or "").strip()
            revised_action = parse_commander_response(rev_text)

            # Option C: revision enforcement — if Commander ignores counter_proposal, escalate to VETO
            cp = oversight_action.counter_proposal
            if cp is not None:
                ignored_action = revised_action.action != cp.action
                ignored_target = revised_action.target_system != cp.target_system
                if ignored_action and ignored_target:
                    print(
                        f"[DEBUG] Revision enforcement: Commander ignored counter_proposal "
                        f"(proposed {revised_action.action}@{revised_action.target_system} "
                        f"vs cp {cp.action}@{cp.target_system}) → auto-VETO",
                        flush=True,
                    )
                    oversight_action = OversightAction(
                        decision=int(OversightDecision.VETO),
                        risk_tier=oversight_action.risk_tier,
                        weakness=f"[AUTO-VETO] Commander ignored counter_proposal. {oversight_action.weakness}",
                        missing_evidence=list(oversight_action.missing_evidence),
                        counter_proposal=cp,
                        lesson_text=oversight_action.lesson_text,
                        lesson_tags=list(oversight_action.lesson_tags),
                    )

            proposal_action = revised_action
            was_revised = True
            _revision_prompt = rev_msg
            _revision_raw = rev_text
        except Exception as exc:
            print(f"[DEBUG] Commander revision LLM error: {exc}", flush=True)

    # --- 4. Apply to env -----------------------------------------------------
    result = await env.step(proposal_action, oversight_action=oversight_action, was_revised=was_revised)

    # --- 5. Record (if recorder attached) ------------------------------------
    if recorder is not None:
        meta = result.observation.metadata or {}
        recorder.record_step(
            step=step_idx + 1,
            hour=int(commander_obs.get("hour", 0)),
            commander_prompt=_commander_prompt_for_recorder,
            commander_raw=_commander_raw_for_recorder,
            commander_parsed={
                "action": proposal_action.action if not was_revised else None,
                "target_system": proposal_action.target_system if not was_revised else None,
                "justification": proposal_action.justification if not was_revised else None,
                "cited_lessons": list(proposal_action.cited_lessons) if not was_revised else None,
                "severity_arg": proposal_action.severity_arg,
                "channel_arg": proposal_action.channel_arg,
                "message_arg": proposal_action.message_arg,
                "scope_arg": proposal_action.scope_arg,
                "evidence_arg": proposal_action.evidence_arg,
                "action_name_initial": ACTION_NAMES.get(proposal_action.action, str(proposal_action.action)),
            },
            oversight_prompt=_oversight_prompt_for_recorder,
            oversight_raw=_oversight_raw_for_recorder,
            oversight_parsed={
                "decision": oversight_action.decision,
                "decision_name": OversightDecision(oversight_action.decision).name,
                "risk_tier": oversight_action.risk_tier,
                "weakness": oversight_action.weakness,
                "missing_evidence": list(oversight_action.missing_evidence),
                "counter_proposal": (
                    oversight_action.counter_proposal.model_dump()
                    if oversight_action.counter_proposal else None
                ),
                "lesson_text": oversight_action.lesson_text,
                "lesson_tags": list(oversight_action.lesson_tags),
            },
            revision_prompt=_revision_prompt,
            revision_raw=_revision_raw,
            revision_parsed=(
                {
                    "action": proposal_action.action,
                    "target_system": proposal_action.target_system,
                    "justification": proposal_action.justification,
                    "cited_lessons": list(proposal_action.cited_lessons),
                    "action_name_final": ACTION_NAMES.get(proposal_action.action, str(proposal_action.action)),
                }
                if was_revised else None
            ),
            env_info=meta,
            commander_reward=float(result.reward or 0.0),
            oversight_reward=float(meta.get("oversight_reward", 0.0)),
            trust_after=meta.get("trust_snapshot") or meta.get("trust_final") or {},
            # Rich dashboard context — all sourced from env metadata
            team_messages=list(result.observation.team_messages or [])
                if hasattr(result.observation, "team_messages") else [],
            siem_alerts=meta.get("step_alerts", []),
            systems_state=meta.get("systems_snapshot", {}),
            investor_state={
                "anxiety": meta.get("investor_anxiety"),
                "tier": meta.get("investor_tier"),
                "persona": meta.get("investor_persona"),
            },
            investor_messages=meta.get("investor_step_messages", []),
            stakeholder_asks=meta.get("stakeholder_new_asks", []),
            governance_events=_extract_governance_events(meta),
            playbook_snapshot=meta.get("playbook_snapshot", []),
            data_exfiltrated=meta.get("data_exfiltrated"),
            stamina=meta.get("team_stamina"),
        )

    return result, proposal_action, oversight_action, was_revised


# ---------------------------------------------------------------------------
# Run one task
# ---------------------------------------------------------------------------

async def run_task(
    env,
    task_id: str,
    client: OpenAI,
    run_root: Optional["Path"] = None,
) -> Dict[str, Any]:
    history: List[str] = []
    oversight_history: List[dict] = []  # Option A: Oversight's own episode memory
    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False
    final_metadata: Dict[str, Any] = {}
    termination_reason = "not_terminated"

    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)

    recorder: Optional[RunRecorder] = None
    try:
        result = await env.reset(task_id=task_id)
        obs = result.observation.model_dump()

        if run_root is not None:
            recorder = RunRecorder(
                run_root=run_root,
                task_id=task_id,
                model_name=MODEL_NAME,
                adversary_gen=obs.get("adversary_gen"),
            )

        for step in range(1, MAX_STEPS + 1):
            if result.done:
                break

            result, final_action, oversight_action, was_revised = await council_step(
                client, env, obs, history, step - 1, recorder=recorder,
                oversight_history=oversight_history,
            )
            obs = result.observation.model_dump()
            reward = result.reward or 0.0
            done = result.done
            err = None
            meta = result.observation.metadata or {}

            action_name = ACTION_NAMES.get(final_action.action, str(final_action.action))
            tgt_name = (
                SYSTEM_NAMES[final_action.target_system]
                if 0 <= final_action.target_system < len(SYSTEM_NAMES) else "?"
            )
            ov_dec = OversightDecision(oversight_action.decision).name
            action_label = f"{action_name}({tgt_name})[{ov_dec}{'+revised' if was_revised else ''}]"

            rewards.append(reward)
            steps_taken = step

            log_step(step=step, action=action_label, reward=reward, done=done, error=err)
            history.append(f"Hour {step}: {action_label} -> {reward:+.2f}")

            # Option A: record Oversight's decision + outcome for next step's episode history
            oversight_history.append({
                "hour": obs.get("hour", step),
                "decision": ov_dec + ("+revised" if was_revised else ""),
                "action_name": action_name,
                "target": tgt_name,
                "outcome": f"r={reward:+.2f} {'DONE' if done else ''}".strip(),
            })

            if done:
                final_scores = meta.get("final_scores") or {}
                score = final_scores.get("final_score", meta.get("comparison_score", 0.5))
                score = min(max(float(score), 0.0), 1.0)
                success = score >= 0.5
                final_metadata = meta
                termination_reason = meta.get("termination_reason", "")
                break

        if not result.done:
            score = 0.5
            success = True
            final_metadata = result.observation.metadata or {}
            termination_reason = "max_steps_without_done"

    except Exception as exc:
        print(f"[DEBUG] Task {task_id} error: {exc}", flush=True)
        score = 0.0
        success = False

    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
        if recorder is not None:
            try:
                paths = recorder.finalize(final_metadata, score=score, success=success)
                print(f"[DEBUG] transcript saved: {paths['json']}", flush=True)
            except Exception as e:
                print(f"[DEBUG] recorder finalize failed: {e}", flush=True)

    return {
        "task_id": task_id,
        "score": score,
        "success": success,
        "steps": steps_taken,
        "duration_s": (recorder.start_ts and (time.time() - recorder.start_ts)) if recorder else 0.0,
        "termination": termination_reason,
        "adversary_gen": final_metadata.get("adversary_gen"),
    }


# ---------------------------------------------------------------------------
# Environment creation (Docker → HF Space → Local fallback)
# ---------------------------------------------------------------------------

async def create_env():
    """Try Docker image first, then HF Space, then local environment."""

    if LOCAL_IMAGE_NAME:
        try:
            from client import CitadelEnv
            print(f"[DEBUG] Trying Docker image: {LOCAL_IMAGE_NAME}", flush=True)
            env = await CitadelEnv.from_docker_image(LOCAL_IMAGE_NAME)
            print("[DEBUG] Docker environment connected", flush=True)
            return env
        except Exception as e:
            print(f"[DEBUG] Docker failed: {e}", flush=True)

    hf_space_url = os.getenv("HF_SPACE_URL", "https://astro-dude-citadel.hf.space")
    try:
        from client import CitadelEnv
        print(f"[DEBUG] Trying HF Space: {hf_space_url}", flush=True)
        env = CitadelEnv(base_url=hf_space_url)
        await env.connect()
        print("[DEBUG] HF Space environment connected", flush=True)
        return env
    except Exception as e:
        print(f"[DEBUG] HF Space failed: {e}", flush=True)

    print("[DEBUG] Using local environment", flush=True)
    return LocalEnv()


def _make_local_env_with_investor(client: Any) -> "LocalEnv":
    """Create a LocalEnv with the shared LLM client wired to the investor agent."""
    return LocalEnv(investor_llm_client=client, investor_model_name=MODEL_NAME)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

async def main() -> None:
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
    env = await create_env()
    # If we ended up with a local env, wire the investor agent to the same LLM client
    if isinstance(env, LocalEnv):
        env = _make_local_env_with_investor(client)

    # Create a timestamped run directory for transcripts (unless disabled)
    run_root = None
    if os.getenv("CITADEL_DISABLE_RECORDING", "").lower() not in ("1", "true", "yes"):
        label = (os.getenv("CITADEL_RUN_LABEL") or MODEL_NAME.replace("/", "-").replace(":", "-"))[:64]
        run_root = make_run_root(label=label)
        print(f"[DEBUG] recording run to {run_root}", flush=True)

    task_results: List[Dict[str, Any]] = []
    try:
        for task_id in TASKS:
            res = await run_task(env, task_id, client, run_root=run_root)
            task_results.append(res)
    finally:
        try:
            await env.close()
        except Exception as e:
            print(f"[DEBUG] env.close() error: {e}", flush=True)

    if run_root is not None and task_results:
        try:
            summary_path = write_run_index(run_root, task_results, MODEL_NAME)
            print(f"[DEBUG] run summary saved: {summary_path}", flush=True)
        except Exception as e:
            print(f"[DEBUG] write_run_index failed: {e}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())