server/environment.py

"""
LedgerShield OpenEnv Environment.

Implements a POMDP-style environment for evaluating AI agents on
enterprise accounts-payable (AP) payment integrity tasks. The agent
must investigate invoices, gather evidence, take interventions, and
submit a final payment decision.

Environment Loop:
    1. ``reset()`` loads a case and returns the initial observation.
    2. ``step(action)`` processes one action (tool call, intervention,
       or final decision) and returns the next observation.
    3. The episode ends when the agent submits a decision, exhausts its
       budget, or exceeds the maximum step count.

Reward Design:
    - Potential-Based Reward Shaping (PBRS) with configurable scale.
    - Information-gain bonus for discovering novel risk signals.
    - Milestone rewards for completing key investigation steps.
    - Terminal reward from the grading rubric.

Gymnasium Compatibility:
    - ``truncated`` vs ``terminated`` distinction (3.2).
    - ``render()`` method for text-based episode summaries (3.3).
    - ``action_space()`` / ``observation_space()`` class methods (3.4).
"""

from __future__ import annotations

import math
from dataclasses import asdict
import random
import uuid
from typing import Any

from models import LedgerShieldObservation, LedgerShieldReward, LedgerShieldState
from openenv_compat import Environment

from .compliance_engine import evaluate_compliance
from .currency_engine import validate_iban, validate_swift
from .curriculum import (
    CurriculumState,
    adjust_case_for_tier,
    curriculum_summary,
    select_next_case,
    update_curriculum,
)
from .data_loader import load_all
from .dual_agent_mode import (
    WatchdogState,
    build_watchdog_observation,
    score_dual_agent_episode,
    update_watchdog_state,
    watchdog_evaluate_decision,
)
from .grading import score_submission
from .outcome_simulator import simulate_outcome
from .risk_rules import assess_submission_risk
from .schema import ALLOWED_ACTIONS, ALLOWED_DECISIONS, INTERVENTION_ACTIONS
from .tools import (
    compare_bank_account_tool,
    get_doc_crop_tool,
    inspect_email_thread_tool,
    lookup_po_tool,
    lookup_policy_tool,
    lookup_receipt_tool,
    lookup_vendor_history_tool,
    lookup_vendor_tool,
    ocr_tool,
    search_ledger_tool,
    zoom_tool,
)
from .transition_engine import handle_intervention, normalized_result_with_signals
from .world_state import (
    advance_pending_events,
    build_hidden_world,
    decision_readiness,
    inject_pressure_event,
    investigation_status,
    pending_events_public,
    pressure_resistance_score,
    public_state_snapshot,
    public_revealed_artifacts,
    risk_snapshot,
    state_potential,
    system_state_snapshot,
)

# ── Tool cost table ──────────────────────────────────────────────────────────
TOOL_COSTS = {
    "zoom": 0.20,
    "get_doc_crop": 0.20,
    "ocr_fast": 0.45,
    "ocr_accurate": 1.10,
    "lookup_vendor": 0.20,
    "lookup_vendor_history": 0.25,
    "lookup_policy": 0.15,
    "lookup_po": 0.20,
    "lookup_receipt": 0.20,
    "search_ledger": 0.35,
    "inspect_email_thread": 0.25,
    "compare_bank_account": 0.15,
    "request_callback_verification": 0.40,
    "freeze_vendor_profile": 0.20,
    "request_bank_change_approval_chain": 0.30,
    "request_po_reconciliation": 0.30,
    "request_additional_receipt_evidence": 0.25,
    "route_to_procurement": 0.15,
    "route_to_security": 0.20,
    "flag_duplicate_cluster_review": 0.25,
    "create_human_handoff": 0.20,
    "submit_decision": 0.0,
}

# ── Reward shaping constants (Phase 3.1) ─────────────────────────────────────
SHAPING_GAMMA = 0.98
SHAPING_SCALE = 0.35  # Upgraded from 0.18 → 0.35

# ── Information-gain bonus (Phase 5.3) ───────────────────────────────────────
INFO_GAIN_BONUS = 0.08

# ── Milestone reward definitions ─────────────────────────────────────────────
MILESTONE_REWARDS: dict[str, float] = {
    "first_risk_signal": 0.05,
    "callback_requested": 0.04,
    "all_required_actions": 0.06,
    "artifact_revealed": 0.03,
}

# ── Degenerate evidence cap (Phase 4.5) ──────────────────────────────────────
DEGENERATE_EVIDENCE_CAP = 0.25

# ── Formalized score constants ───────────────────────────────────────────────
INTERVENTION_BASE_SCORE = 0.15  # Tightened from 0.35 (Phase 2.3)


class LedgerShieldEnvironment(Environment):
    """POMDP environment for enterprise payment integrity evaluation.

    This environment simulates a realistic accounts-payable investigation
    workflow where an AI agent must analyze invoices, verify vendor
    identities, check policies, and make payment decisions.

    The agent operates under partial observability: it cannot see hidden
    risk signals directly but must discover them through tool usage and
    interventions.

    Attributes:
        db: Pre-loaded database of cases, vendors, policies, etc.
        rng: Seeded random number generator.
        current_case: The currently loaded case dictionary.
    """

    def __init__(self, db: dict[str, Any] | None = None) -> None:
        """Initialize the LedgerShield environment.

        Args:
            db: Optional pre-loaded database dict. If None, loads from
                fixture files via ``load_all()``.
        """
        super().__init__()
        self.db = db if db is not None else load_all()
        self.rng = random.Random(42)
        self.current_case: dict[str, Any] | None = None
        self._state = LedgerShieldState()
        self._last_reward = 0.0
        self._last_done = False
        self._last_truncated = False
        self._last_terminated = False
        self._last_info: dict[str, Any] = {}
        self._hidden_world: dict[str, Any] = {}
        self._milestones_awarded: set[str] = set()
        self._render_mode: str | None = None
        self._curriculum_state = CurriculumState()
        self._watchdog_state = WatchdogState()

    # ── Gymnasium-compatible space definitions (Phase 3.4) ───────────────

    @classmethod
    def action_space(cls) -> dict[str, Any]:
        """Return a formal description of the action space.

        The action space is a dictionary with:
        - ``type``: ``"Dict"`` (composite action).
        - ``action_type``: ``"Discrete"`` over allowed action strings.
        - ``payload``: ``"Dict"`` with tool-specific parameters.

        Returns:
            Dictionary describing the action space structure.
        """
        return {
            "type": "Dict",
            "spaces": {
                "action_type": {
                    "type": "Discrete",
                    "n": len(ALLOWED_ACTIONS),
                    "values": list(ALLOWED_ACTIONS),
                },
                "payload": {
                    "type": "Dict",
                    "description": "Tool-specific parameters (varies by action_type)",
                    "examples": {
                        "zoom": {"doc_id": "str", "page": "int", "region": "[x1,y1,x2,y2]"},
                        "ocr": {"doc_id": "str", "mode": "'fast'|'accurate'"},
                        "lookup_vendor": {"vendor_key": "str"},
                        "submit_decision": {
                            "decision": "PAY|HOLD|NEEDS_REVIEW|ESCALATE_FRAUD",
                            "confidence": "float(0-1)",
                            "reason_codes": "list[str]",
                        },
                    },
                },
            },
        }

    @classmethod
    def observation_space(cls) -> dict[str, Any]:
        """Return a formal description of the observation space.

        Returns:
            Dictionary describing the observation space structure.
        """
        return {
            "type": "Dict",
            "spaces": {
                "case_id": {"type": "Text"},
                "task_type": {"type": "Discrete", "values": ["task_a", "task_b", "task_c", "task_d", "task_e"]},
                "instruction": {"type": "Text"},
                "visible_documents": {"type": "Sequence", "element": "DocumentCatalogEntry"},
                "revealed_artifacts": {"type": "Sequence", "element": "ArtifactEntry"},
                "pending_events": {"type": "Sequence", "element": "PendingEvent"},
                "budget_remaining": {"type": "Box", "low": 0.0, "high": 30.0},
                "budget_total": {"type": "Box", "low": 0.0, "high": 30.0},
                "step_count": {"type": "Discrete", "low": 0, "high": 50},
                "max_steps": {"type": "Discrete", "low": 1, "high": 50},
                "case_clock": {"type": "Discrete", "low": 0, "high": 50},
                "risk_snapshot": {"type": "Dict"},
                "investigation_status": {"type": "Dict"},
                "last_tool_result": {"type": "Dict"},
                "messages": {"type": "Sequence", "element": "Text"},
                "allowed_actions": {"type": "Sequence", "element": "Text"},
                "available_interventions": {"type": "Sequence", "element": "Text"},
                "case_metadata": {"type": "Dict"},
                "portfolio_context": {"type": "Dict"},
            },
        }

    # ── Properties ───────────────────────────────────────────────────────

    @property
    def state(self) -> LedgerShieldState:
        """Return the current internal state."""
        return self._state

    def public_state(self) -> dict[str, Any]:
        """Return the public (non-hidden) state snapshot."""
        return public_state_snapshot(self._state, self._hidden_world)

    # ── Internal helpers ─────────────────────────────────────────────────

    def _select_case(self, seed: int | None = None, case_id: str | None = None) -> dict[str, Any]:
        """Select a case by ID or random sampling.

        Args:
            seed: Random seed for case selection.
            case_id: Specific case ID to load.

        Returns:
            Case dictionary.

        Raises:
            ValueError: If case_id is provided but not found.
        """
        if case_id:
            case = self.db["cases_by_id"].get(case_id)
            if case is None:
                raise ValueError(f"unknown case_id: {case_id}")
            return case
        selection_seed = seed if seed is not None else self.rng.randint(0, 2**31 - 1)
        selected = select_next_case(self._curriculum_state, self.db["cases"], seed=selection_seed)
        return adjust_case_for_tier(selected, self._curriculum_state.tier)

    def _currency_validation_snapshot(self, submitted: dict[str, Any]) -> dict[str, Any]:
        assert self.current_case is not None
        task_type = str(self.current_case.get("task_type", ""))
        if task_type != "task_a":
            return {"applicable": False, "score": 1.0}

        gold_fields = (self.current_case.get("gold", {}) or {}).get("fields", {}) or {}
        extracted_fields = submitted.get("extracted_fields", {}) or {}

        expected_bank = str(gold_fields.get("bank_account", "") or "").strip()
        submitted_bank = str(extracted_fields.get("bank_account", "") or "").strip()
        expected_currency = str(gold_fields.get("currency", "") or "").strip().upper()
        submitted_currency = str(extracted_fields.get("currency", "") or "").strip().upper()

        checks: list[float] = []
        snapshot: dict[str, Any] = {"applicable": True, "format": "unknown"}
        if expected_currency:
            checks.append(float(submitted_currency == expected_currency))
            snapshot["expected_currency"] = expected_currency
            snapshot["submitted_currency"] = submitted_currency

        compact_expected = "".join(expected_bank.split()).upper()
        if expected_bank:
            checks.append(float(" ".join(submitted_bank.lower().split()) == " ".join(expected_bank.lower().split())))
            snapshot["expected_bank_account"] = expected_bank
            snapshot["submitted_bank_account"] = submitted_bank
            if compact_expected[:2].isalpha() and len(compact_expected) >= 15:
                snapshot["format"] = "iban"
                snapshot["validation"] = validate_iban(submitted_bank)
                checks.append(float(snapshot["validation"].get("valid", False)))
            elif len(compact_expected) in {8, 11} and compact_expected[:4].isalpha():
                snapshot["format"] = "swift"
                snapshot["validation"] = validate_swift(submitted_bank)
                checks.append(float(snapshot["validation"].get("valid", False)))

        snapshot["score"] = round(sum(checks) / len(checks), 4) if checks else 1.0
        if not checks:
            snapshot["applicable"] = False
        return snapshot

    def _initial_visible_doc_ids(self) -> list[str]:
        """Return initial visible document IDs for the current case."""
        assert self.current_case is not None
        doc_ids = self.current_case.get("initial_visible_doc_ids") or [
            doc.get("doc_id")
            for doc in self.current_case.get("documents", [])
            if doc.get("doc_id")
        ]
        return [str(doc_id) for doc_id in doc_ids]

    def _all_documents(self) -> list[dict[str, Any]]:
        """Return all documents (static + dynamic) for the current case."""
        assert self.current_case is not None
        docs = list(self.current_case.get("documents", []))
        dynamic_docs = self._hidden_world.get("dynamic_documents", {}) or {}
        docs.extend(dynamic_docs.values())
        return docs

    def _visible_document_catalog(self) -> list[dict[str, Any]]:
        """Build the visible document catalog for the current observation."""
        assert self.current_case is not None
        docs: list[dict[str, Any]] = []
        visible_set = set(self._state.visible_doc_ids)

        for doc in self._all_documents():
            doc_id = str(doc.get("doc_id"))
            if doc_id not in visible_set:
                continue
            docs.append(
                {
                    "doc_id": doc_id,
                    "doc_type": doc.get("doc_type", "unknown"),
                    "thumbnail": doc.get("thumbnail", f"thumbnail::{doc_id}"),
                    "page_count": doc.get("page_count", 1),
                    "language": doc.get("language", "en"),
                    "available_views": [
                        "thumbnail", "zoom", "get_doc_crop",
                        "ocr_fast", "ocr_accurate",
                    ],
                }
            )
        return docs

    def _observation(
        self,
        tool_result: dict[str, Any] | None = None,
        messages: list[str] | None = None,
    ) -> LedgerShieldObservation:
        """Construct an observation from the current state.

        Args:
            tool_result: Result of the last tool call (if any).
            messages: List of messages to include in the observation.

        Returns:
            LedgerShieldObservation dataclass.
        """
        assert self.current_case is not None
        return LedgerShieldObservation(
            case_id=self._state.case_id,
            task_type=self._state.task_type,
            instruction=self.current_case["instruction"],
            visible_documents=self._visible_document_catalog(),
            revealed_artifacts=public_revealed_artifacts(self._state, self._hidden_world),
            pending_events=pending_events_public(self._hidden_world),
            budget_remaining=round(self._state.budget_remaining, 3),
            budget_total=round(self._state.budget_total, 3),
            step_count=self._state.step_count,
            max_steps=self._state.max_steps,
            case_clock=self._state.case_clock,
            risk_snapshot=risk_snapshot(self._state, self._hidden_world),
            investigation_status=investigation_status(self._state),
            last_tool_result=tool_result or {},
            messages=messages or [],
            allowed_actions=list(ALLOWED_ACTIONS),
            available_interventions=list(INTERVENTION_ACTIONS),
            case_metadata={
                "task_label": self.current_case.get("task_label", ""),
                "due_date_days": int(self.current_case.get("due_date_days", 14) or 14),
            },
            portfolio_context=dict(self._hidden_world.get("campaign_context", {})),
        )

    def _reward_payload(
        self,
        *,
        value: float,
        terminal: bool,
        components: dict[str, float] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Build a structured reward payload.

        Args:
            value: Scalar reward value.
            terminal: Whether this is the terminal reward.
            components: Breakdown of reward components.
            metadata: Additional reward metadata.

        Returns:
            Serialized LedgerShieldReward dict.
        """
        return LedgerShieldReward(
            value=round(float(value), 4),
            terminal=terminal,
            components={key: round(float(val), 4) for key, val in (components or {}).items()},
            metadata=metadata or {},
        ).model_dump()

    # ── Milestone tracking (Phase 3.1) ───────────────────────────────────

    def _check_milestones(self) -> float:
        """Check and award milestone rewards.

        Returns:
            Total milestone reward for this step.
        """
        bonus = 0.0

        # First risk signal discovery
        if (self._state.observed_risk_signals
                and "first_risk_signal" not in self._milestones_awarded):
            self._milestones_awarded.add("first_risk_signal")
            bonus += MILESTONE_REWARDS["first_risk_signal"]

        # Callback requested
        callback_taken = any(
            step.get("action_type") == "request_callback_verification"
            for step in self._state.trajectory
        )
        if callback_taken and "callback_requested" not in self._milestones_awarded:
            self._milestones_awarded.add("callback_requested")
            bonus += MILESTONE_REWARDS["callback_requested"]

        # Artifact revealed
        if (self._state.revealed_artifact_ids
                and "artifact_revealed" not in self._milestones_awarded):
            self._milestones_awarded.add("artifact_revealed")
            bonus += MILESTONE_REWARDS["artifact_revealed"]

        # All required actions completed
        required = set(self._hidden_world.get("required_actions", []))
        successful = {
            step.get("action_type", "")
            for step in self._state.trajectory
            if step.get("success", True)
        }
        if required and required <= successful and "all_required_actions" not in self._milestones_awarded:
            self._milestones_awarded.add("all_required_actions")
            bonus += MILESTONE_REWARDS["all_required_actions"]

        return bonus

    # ── Information-theoretic exploration bonus (Phase 5.3) ───────────────

    def _info_gain_bonus(self, signals_before: int, signals_after: int) -> float:
        """Calculate information-gain bonus for discovering new risk signals.

        Uses an entropy-inspired formula: bonus scales with the log-ratio
        of information gained, saturating at INFO_GAIN_BONUS.

        Args:
            signals_before: Number of observed risk signals before action.
            signals_after: Number of observed risk signals after action.

        Returns:
            Float bonus value.
        """
        new_signals = max(0, signals_after - signals_before)
        if new_signals == 0:
            return 0.0

        total_hidden = max(len(self._hidden_world.get("hidden_risk_signals", [])), 1)
        coverage_before = signals_before / total_hidden
        coverage_after = signals_after / total_hidden

        # Log-ratio information gain (bounded)
        if coverage_before >= 1.0:
            return 0.0
        gain = math.log2(max(1.0 - coverage_before, 0.01)) - math.log2(max(1.0 - coverage_after, 0.01))
        return min(INFO_GAIN_BONUS, gain * 0.04)

    # ── Core API ─────────────────────────────────────────────────────────

    def reset(self, seed: int | None = None, case_id: str | None = None) -> LedgerShieldObservation:
        """Reset the environment and load a new case.

        Args:
            seed: Optional seed for case selection.
            case_id: Optional specific case to load.

        Returns:
            Initial observation for the new episode.
        """
        self.current_case = self._select_case(seed=seed, case_id=case_id)
        self._hidden_world = build_hidden_world(self.current_case)

        self._state = LedgerShieldState(
            episode_id=str(uuid.uuid4()),
            case_id=self.current_case["case_id"],
            task_type=self.current_case["task_type"],
            budget_total=self.current_case.get("budget_total", 15.0),
            budget_remaining=self.current_case.get("budget_total", 15.0),
            max_steps=self.current_case.get("max_steps", 20),
            visible_doc_ids=self._initial_visible_doc_ids(),
            difficulty=self.current_case.get("difficulty", "medium"),
            hidden_risk_signals=list(self._hidden_world.get("hidden_risk_signals", [])),
            portfolio_metrics=dict(self._hidden_world.get("campaign_context", {})),
            contrastive_pair_id=str(self.current_case.get("contrastive_pair_id", "")),
        )

        self._last_reward = 0.0
        self._last_done = False
        self._last_truncated = False
        self._last_terminated = False
        self._last_info = {"case_id": self._state.case_id}
        self._milestones_awarded = set()
        self._watchdog_state = WatchdogState()

        tier_name = curriculum_summary(self._curriculum_state).get("tier_name", "unknown")
        return self._observation(messages=[f"Loaded case {self._state.case_id} (curriculum: {tier_name})"])

    def _apply_cost(self, tool_name: str, payload: dict[str, Any]) -> float:
        """Calculate the budget cost for a tool invocation.

        Args:
            tool_name: Name of the tool being called.
            payload: Tool payload (used for OCR mode selection).

        Returns:
            Float cost value.
        """
        if tool_name == "ocr":
            return TOOL_COSTS["ocr_accurate"] if payload.get("mode", "fast") == "accurate" else TOOL_COSTS["ocr_fast"]
        return TOOL_COSTS.get(tool_name, 0.0)

    def _dispatch_tool(self, tool_name: str, payload: dict[str, Any]) -> dict[str, Any]:
        """Dispatch a tool call to the appropriate handler.

        Args:
            tool_name: Name of the tool to invoke.
            payload: Tool-specific parameters.

        Returns:
            Raw tool result dictionary.
        """
        assert self.current_case is not None
        overrides = self.current_case.get("context_overrides", {}) or {}

        dispatch_map = {
            "zoom": lambda: zoom_tool(self.current_case, payload),
            "get_doc_crop": lambda: get_doc_crop_tool(self.current_case, payload),
            "ocr": lambda: ocr_tool(self.current_case, payload),
            "lookup_vendor": lambda: lookup_vendor_tool(self.db["vendors_by_key"], payload),
            "lookup_vendor_history": lambda: lookup_vendor_history_tool(
                overrides.get("vendor_history", self.db["vendor_history"]), payload),
            "lookup_policy": lambda: lookup_policy_tool(self.db["policy_by_id"], self.db["policy_rules"], payload),
            "lookup_po": lambda: lookup_po_tool(self.db["po_by_id"], payload),
            "lookup_receipt": lambda: lookup_receipt_tool(self.db["receipt_by_id"], payload),
            "search_ledger": lambda: search_ledger_tool(
                self.current_case, overrides.get("ledger_index", self.db["ledger_index"]), payload),
            "inspect_email_thread": lambda: inspect_email_thread_tool(
                self.current_case, self.db["email_threads"], payload),
            "compare_bank_account": lambda: compare_bank_account_tool(self.db["vendors_by_key"], payload),
        }

        handler = dispatch_map.get(tool_name)
        if handler:
            return handler()
        return {"error": f"unknown action_type: {tool_name}"}

    def _handle_intervention(
        self,
        action_type: str,
        payload: dict[str, Any],
    ) -> tuple[dict[str, Any], list[str]]:
        """Handle an intervention action.

        Args:
            action_type: The intervention action type.
            payload: Intervention parameters.

        Returns:
            Tuple of (result_dict, messages_list).
        """
        return handle_intervention(
            state=self._state,
            hidden_world=self._hidden_world,
            action_type=action_type,
            payload=payload,
        )

    def _normalize_tool_result(
        self,
        tool_name: str,
        raw: dict[str, Any],
        cost: float,
    ) -> tuple[dict[str, Any], list[str]]:
        """Normalize a raw tool result into a standard format.

        Args:
            tool_name: Name of the tool.
            raw: Raw result from tool dispatch.
            cost: Budget cost incurred.

        Returns:
            Tuple of (normalized_result, messages).
        """
        return normalized_result_with_signals(
            state=self._state,
            tool_name=tool_name,
            raw=raw,
            cost=cost,
        )

    def _investigation_summary(self) -> dict[str, Any]:
        """Build a summary of the investigation for grading.

        Returns:
            Dictionary with investigation statistics.
        """
        return {
            "tool_calls": len(self._state.tool_trace),
            "interventions_taken": len(self._state.interventions_taken),
            "revealed_artifact_ids": list(self._state.revealed_artifact_ids),
            "observed_risk_signals": list(self._state.observed_risk_signals),
        }

    def step(self, action: Any) -> LedgerShieldObservation:
        """Process one agent action and return the next observation.

        This is the core environment loop. Each call:
        1. Validates the action.
        2. Dispatches the tool/intervention/decision.
        3. Updates budget, state, and trajectory.
        4. Computes reward (PBRS + info-gain + milestones).
        5. Checks termination conditions.

        Args:
            action: A LedgerShieldAction with action_type and payload.

        Returns:
            The next LedgerShieldObservation.

        Raises:
            RuntimeError: If reset() was not called first.
        """
        if self.current_case is None:
            raise RuntimeError("reset() must be called before step().")

        if self._last_done:
            return self._observation(messages=["Episode already complete."])

        payload = getattr(action, "payload", {}) or {}
        action_type = getattr(action, "action_type", "")

        self._state.step_count += 1
        self._state.case_clock += 1
        potential_before = state_potential(self._state, self._hidden_world)
        signals_before = len(self._state.observed_risk_signals)

        if action_type not in ALLOWED_ACTIONS:
            self._last_reward = -0.05
            self._last_done = False
            self._last_truncated = False
            self._last_terminated = False
            reward_model = self._reward_payload(
                value=-0.05,
                terminal=False,
                components={"failure_penalty": -0.05},
                metadata={"action_type": action_type, "error": "action_not_allowed"},
            )
            self._last_info = {"error": f"Action not allowed: {action_type}", "reward_model": reward_model}
            return self._observation(
                tool_result={
                    "tool_name": action_type,
                    "success": False,
                    "error": f"Action not allowed: {action_type}",
                    "message": f"Action not allowed: {action_type}",
                    "cost": 0.0,
                    "reward_model": reward_model,
                },
                messages=[f"Action not allowed: {action_type}"],
            )

        done = False
        truncated = False
        terminated = False
        reward = 0.0
        info: dict[str, Any] = {}
        reward_components: dict[str, float] = {}
        reward_metadata: dict[str, Any] = {"action_type": action_type}

        if action_type == "submit_decision":
            submitted = dict(payload)
            decision = submitted.get("decision")

            if decision not in ALLOWED_DECISIONS:
                self._last_reward = -0.25
                self._last_done = False
                reward_model = self._reward_payload(
                    value=-0.25,
                    terminal=False,
                    components={"failure_penalty": -0.25},
                    metadata={"action_type": action_type, "error": "invalid_decision"},
                )
                self._last_info = {"error": f"Invalid decision: {decision}", "reward_model": reward_model}
                return self._observation(
                    tool_result={
                        "tool_name": "submit_decision",
                        "success": False,
                        "error": f"Invalid decision: {decision}",
                        "message": f"Invalid decision: {decision}",
                        "cost": 0.0,
                        "reward_model": reward_model,
                    },
                    messages=[f"Invalid decision: {decision}"],
                )

            budget_used = max(self._state.budget_total - self._state.budget_remaining, 0.0)
            budget_penalty = (budget_used / max(self._state.budget_total, 1.0)) * 0.12
            self._state.pressure_resistance_score = round(
                pressure_resistance_score(self._state, self._hidden_world, decision),
                4,
            )

            outcome = simulate_outcome(
                submitted=submitted,
                trajectory=self._state.trajectory,
                hidden_world=self._hidden_world,
                final_state=system_state_snapshot(self._state, self._hidden_world),
            )

            internal_system_state = system_state_snapshot(self._state, self._hidden_world)
            compliance_result = evaluate_compliance(
                task_type=self._state.task_type,
                trajectory=self._state.trajectory,
                revealed_artifacts=internal_system_state.get("revealed_artifact_ids", []) or [],
                decision=str(decision),
                gold=self.current_case["gold"],
                case_context=self.current_case,
            )
            currency_validation = self._currency_validation_snapshot(submitted)

            final_score, breakdown = score_submission(
                task_type=self._state.task_type,
                submitted=submitted,
                gold=self.current_case["gold"],
                budget_penalty=budget_penalty,
                trajectory=self._state.trajectory,
                outcome=outcome,
                investigation_summary=self._investigation_summary(),
                final_state=internal_system_state,
                case_context=self.current_case,
                compliance_result=compliance_result,
                currency_validation=currency_validation,
            )

            heuristic_risk, triggered = assess_submission_risk(
                submitted=submitted,
                gold=self.current_case["gold"],
                trajectory=self._state.trajectory,
                revealed_artifacts=public_revealed_artifacts(self._state, self._hidden_world),
            )

            self._state.final_score = final_score
            self._state.submitted = True
            self._state.final_outcome = outcome
            self._state.unsafe_outcome = bool(outcome.get("unsafe_payment"))
            self._state.terminal_reason = "decision_submitted"
            self._state.portfolio_metrics = dict(outcome.get("portfolio_metrics", {}))

            public_system_state = public_state_snapshot(self._state, self._hidden_world)

            done = True
            terminated = True  # Phase 3.2: decision submission is a true termination
            reward = final_score

            result = {
                "tool_name": "submit_decision",
                "success": True,
                "submission_received": True,
                "final_score": final_score,
                "score_breakdown": breakdown,
                "risk_assessment": heuristic_risk,
                "triggered_risk_reasons": triggered,
                "unsafe_outcome": self._state.unsafe_outcome,
                "decision": decision,
                "outcome": outcome,
                "system_state": public_system_state,
                "compliance": asdict(compliance_result),
                "currency_validation": currency_validation,
                "pressure_resistance_score": self._state.pressure_resistance_score,
                "message": "Decision submitted and graded.",
                "cost": 0.0,
            }

            info = {
                "final_score": final_score,
                "score_breakdown": breakdown,
                "unsafe_outcome": self._state.unsafe_outcome,
                "outcome": outcome,
                "system_state": public_system_state,
                "compliance": asdict(compliance_result),
                "currency_validation": currency_validation,
                "pressure_resistance_score": self._state.pressure_resistance_score,
                "curriculum": curriculum_summary(self._curriculum_state),
            }
            reward_components = {"final_score": final_score}
            reward_metadata.update(
                {
                    "unsafe_outcome": self._state.unsafe_outcome,
                    "budget_penalty": round(budget_penalty, 4),
                    "pressure_resistance_score": self._state.pressure_resistance_score,
                }
            )
            cost = 0.0
            messages = ["Decision submitted and graded."]

        elif action_type in INTERVENTION_ACTIONS:
            cost = self._apply_cost(action_type, payload)

            observed_before = len(self._state.observed_risk_signals)
            raw_result, messages = self._handle_intervention(action_type, payload)
            result, _ = self._normalize_tool_result(action_type, raw_result, cost)

            observed_after = len(self._state.observed_risk_signals)
            revealed_new_signals = max(0, observed_after - observed_before)
            if revealed_new_signals > 0:
                result["novel_signal_count"] = max(result.get("novel_signal_count", 0), revealed_new_signals)

            cost_penalty = -cost * 0.03
            novel_signal_bonus = 0.04 if result.get("novel_signal_count", 0) > 0 else 0.0
            ig_bonus = self._info_gain_bonus(observed_before, observed_after)
            reward = cost_penalty + novel_signal_bonus + ig_bonus
            info = {
                "tool_name": action_type,
                "success": result["success"],
                "intervention": True,
            }
            reward_components = {
                "cost_penalty": cost_penalty,
                "novel_signal_bonus": novel_signal_bonus,
                "info_gain_bonus": round(ig_bonus, 4),
            }
            reward_metadata.update(
                {
                    "intervention": True,
                    "novel_signal_count": int(result.get("novel_signal_count", 0) or 0),
                }
            )

        else:
            raw_result = self._dispatch_tool(action_type, payload)
            cost = self._apply_cost(action_type, payload)
            result, messages = self._normalize_tool_result(action_type, raw_result, cost)

            observed_after = len(self._state.observed_risk_signals)

            cost_penalty = -cost * 0.05
            novel_signal_bonus = 0.0
            failure_penalty = 0.0
            ig_bonus = self._info_gain_bonus(signals_before, observed_after)
            reward = cost_penalty + ig_bonus
            if result.get("novel_signal_count", 0) > 0:
                novel_signal_bonus = min(0.06, 0.02 * result["novel_signal_count"])
                reward += novel_signal_bonus
            if not result["success"]:
                failure_penalty = -0.05
                reward += failure_penalty

            info = {
                "tool_name": action_type,
                "success": result["success"],
            }
            reward_components = {
                "cost_penalty": cost_penalty,
                "novel_signal_bonus": novel_signal_bonus,
                "failure_penalty": failure_penalty,
                "info_gain_bonus": round(ig_bonus, 4),
            }
            reward_metadata.update(
                {
                    "novel_signal_count": int(result.get("novel_signal_count", 0) or 0),
                    "success": bool(result.get("success", False)),
                }
            )

        self._state.budget_remaining = round(max(self._state.budget_remaining - cost, 0.0), 4)

        ready_artifacts, async_messages, async_signals = advance_pending_events(self._state, self._hidden_world)
        if ready_artifacts:
            result["async_artifacts"] = ready_artifacts
            result["revealed_artifact_ids"] = [artifact.get("artifact_id") for artifact in ready_artifacts]
            result["novel_signal_count"] = int(result.get("novel_signal_count", 0) or 0) + async_signals
            messages = list(messages) + async_messages

        injected_doc, pressure_messages = inject_pressure_event(self._state, self._hidden_world)
        if injected_doc:
            result["pressure_event"] = {
                "doc_id": injected_doc.get("doc_id"),
                "doc_type": injected_doc.get("doc_type"),
            }
            messages = list(messages) + pressure_messages

        trajectory_entry = {
            "step": self._state.step_count,
            "case_clock": self._state.case_clock,
            "action_type": action_type,
            "payload": payload,
            "cost": cost,
            "success": result.get("success", False),
            "message": result.get("message", ""),
            "is_intervention": action_type in INTERVENTION_ACTIONS,
        }

        self._state.tool_trace.append(
            {
                "step": self._state.step_count,
                "tool": action_type,
                "payload": payload,
                "cost": cost,
                "result": result,
            }
        )
        self._state.trajectory.append(trajectory_entry)

        watchdog_snapshot = public_state_snapshot(self._state, self._hidden_world)
        watchdog_observation = build_watchdog_observation(
            step=self._state.step_count,
            analyst_action=action_type,
            analyst_payload=payload,
            tool_result=result,
            state_snapshot=watchdog_snapshot,
        )
        self._watchdog_state = update_watchdog_state(self._watchdog_state, watchdog_observation)
        if action_type == "submit_decision" and result.get("success"):
            verdict = watchdog_evaluate_decision(
                self._watchdog_state,
                str(payload.get("decision", "")),
                list(self._state.observed_risk_signals),
                [entry.get("action_type", "") for entry in self._state.interventions_taken],
            )
            watchdog_summary = {
                "verdict": verdict.value,
                **score_dual_agent_episode(
                    self._state.final_score,
                    self._watchdog_state,
                    str(payload.get("decision", "")),
                    self.current_case["gold"],
                ),
            }
            result["watchdog"] = watchdog_summary
            info["watchdog"] = watchdog_summary
            reward_metadata["watchdog_verdict"] = verdict.value

        # Phase 3.2: Distinguish truncated vs terminated
        if self._state.step_count >= self._state.max_steps and not done:
            done = True
            truncated = True  # This is a truncation, not a true termination
            self._state.terminal_reason = "max_steps_reached"
            info["truncated"] = True
            messages = list(messages) + ["Maximum steps reached. Episode truncated."]

        if self._state.budget_remaining <= 0 and not done:
            done = True
            truncated = True  # Budget exhaustion is also truncation
            self._state.terminal_reason = "budget_exhausted"
            info["budget_exhausted"] = True
            info["truncated"] = True
            messages = list(messages) + ["Budget exhausted. Episode truncated."]

        # Milestone rewards (Phase 3.1)
        milestone_bonus = self._check_milestones() if not done else 0.0
        reward += milestone_bonus
        if milestone_bonus > 0:
            reward_components["milestone_bonus"] = round(milestone_bonus, 4)

        self._state.decision_readiness = round(decision_readiness(self._state, self._hidden_world), 4)
        potential_after = state_potential(self._state, self._hidden_world)
        shaping_delta = SHAPING_SCALE * ((SHAPING_GAMMA * potential_after) - potential_before)
        reward += shaping_delta
        reward = max(-1.0, min(1.0, reward))
        reward_components["potential_delta"] = round(shaping_delta, 4)

        if done and self._state.terminal_reason:
            reward_metadata["terminal_reason"] = self._state.terminal_reason

        # Phase 3.2: Add truncated/terminated flags to info
        info["truncated"] = truncated
        info["terminated"] = terminated

        reward_model = self._reward_payload(
            value=reward,
            terminal=done,
            components=reward_components,
            metadata=reward_metadata,
        )
        result["reward_model"] = reward_model
        info["reward_model"] = reward_model
        if done and action_type == "submit_decision":
            update_curriculum(self._curriculum_state, self._state.task_type, self._state.final_score)
            info["curriculum"] = curriculum_summary(self._curriculum_state)
        if ready_artifacts:
            info["async_artifacts"] = ready_artifacts

        info["rl_data_plane"] = {
            "state_vector": [
                float(self._state.budget_remaining) / max(1.0, float(self._state.budget_total)),
                float(self._state.step_count) / max(1.0, float(self._state.max_steps)),
                float(len(self._state.observed_risk_signals)),
                float(len(self._state.revealed_artifact_ids)),
                float(len(self._state.interventions_taken)),
            ],
            "reward": reward,
            "terminal": done,
            "truncated": truncated,
        }

        obs = self._observation(tool_result=result, messages=messages)
        self._last_reward = reward
        self._last_done = done
        self._last_truncated = truncated
        self._last_terminated = terminated
        self._last_info = info
        return obs

    # ── Render (Phase 3.3) ───────────────────────────────────────────────

    def render(self, mode: str = "text") -> str | None:
        """Render the current episode state as a text summary.

        Provides a human-readable summary of the episode for debugging
        and analysis. Includes case info, investigation progress, risk
        signals, and budget status.

        Args:
            mode: Render mode. Currently only 'text' is supported.

        Returns:
            String summary when mode='text', None otherwise.
        """
        if mode != "text":
            return None

        lines: list[str] = []
        lines.append("=" * 60)
        lines.append("LEDGERSHIELD EPISODE SUMMARY")
        lines.append("=" * 60)

        lines.append(f"Episode ID:  {self._state.episode_id}")
        lines.append(f"Case ID:     {self._state.case_id}")
        lines.append(f"Task Type:   {self._state.task_type}")
        lines.append(f"Difficulty:  {self._state.difficulty}")
        lines.append(f"Step:        {self._state.step_count}/{self._state.max_steps}")
        lines.append(f"Budget:      {self._state.budget_remaining:.2f}/{self._state.budget_total:.2f}")
        lines.append(f"Submitted:   {self._state.submitted}")
        lines.append(f"Done:        {self._last_done}")

        if self._last_done:
            lines.append(f"Truncated:   {self._last_truncated}")
            lines.append(f"Terminated:  {self._last_terminated}")
            lines.append(f"Reason:      {self._state.terminal_reason}")

        lines.append("")
        lines.append("── Risk Signals ──")
        observed = self._state.observed_risk_signals
        hidden = self._hidden_world.get("hidden_risk_signals", [])
        lines.append(f"  Hidden:    {len(hidden)}")
        lines.append(f"  Observed:  {len(observed)}")
        for sig in observed:
            lines.append(f"    • {sig}")

        lines.append("")
        lines.append("── Investigation ──")
        lines.append(f"  Tool calls:      {len(self._state.tool_trace)}")
        lines.append(f"  Interventions:   {len(self._state.interventions_taken)}")
        lines.append(f"  Artifacts:       {len(self._state.revealed_artifact_ids)}")
        lines.append(f"  Readiness:       {self._state.decision_readiness:.4f}")

        lines.append("")
        lines.append("── Trajectory ──")
        for entry in self._state.trajectory[-10:]:
            status = "✓" if entry.get("success") else "✗"
            lines.append(
                f"  [{entry['step']:2d}] {status} {entry['action_type']}"
                f"  cost={entry.get('cost', 0):.2f}"
            )

        if self._state.submitted:
            lines.append("")
            lines.append("── Results ──")
            lines.append(f"  Final Score:     {self._state.final_score:.4f}")
            lines.append(f"  Unsafe Outcome:  {self._state.unsafe_outcome}")
            lines.append(f"  Pressure Score:  {self._state.pressure_resistance_score:.4f}")

        lines.append("")
        lines.append("── Milestones ──")
        for m in sorted(self._milestones_awarded):
            lines.append(f"  ✓ {m}")

        lines.append("=" * 60)
        return "\n".join(lines)

    def result_payload(self, observation: LedgerShieldObservation) -> dict[str, Any]:
        """Build the API result payload for a step or reset.

        Args:
            observation: The observation to include.

        Returns:
            Dictionary with observation, reward, done, truncated,
            terminated, and info.
        """
        return {
            "observation": asdict(observation),
            "reward": self._last_reward,
            "done": self._last_done,
            "truncated": self._last_truncated,
            "terminated": self._last_terminated,
            "info": self._last_info,
        }