server/world_state.py

from __future__ import annotations

from copy import deepcopy
from typing import Any
import uuid

from models import LedgerShieldState

from .pressure_events import PressureEvent, evaluate_pressure_resistance, schedule_pressure_event
from .risk_rules import compute_due_date_potential, derive_case_risk_signals, risk_bucket
from .schema import canonical_reason_codes, normalize_text
from .vendor_simulator import build_vendor_simulator_state, simulate_callback
from .evidence_graph import EvidenceGraph, generate_scenario_graph


def _successful_actions(state: LedgerShieldState) -> set[str]:
    return {
        normalize_text(step.get("action_type"))
        for step in state.trajectory
        if step.get("success", True)
    }


def _required_actions(case: dict[str, Any], hidden_signals: list[str]) -> list[str]:
    task_type = normalize_text(case.get("task_type"))
    required = {
        "task_a": ["ocr", "zoom"],
        "task_b": ["lookup_policy", "lookup_po", "lookup_receipt"],
        "task_c": ["search_ledger", "compare_bank_account"],
        "task_d": ["inspect_email_thread", "lookup_vendor_history", "lookup_policy", "compare_bank_account", "search_ledger"],
        "task_e": [
            "inspect_email_thread",
            "lookup_vendor_history",
            "lookup_policy",
            "compare_bank_account",
            "search_ledger",
            "request_callback_verification",
            "flag_duplicate_cluster_review",
            "route_to_security",
            "freeze_vendor_profile",
        ],
    }.get(task_type, [])
    hidden = {normalize_text(signal) for signal in hidden_signals}

    seen: set[str] = set()
    ordered: list[str] = []
    
    # Extract dynamically required actions from the EvidenceGraph if provided
    use_legacy_rules = True
    latent_graph_data = case.get("graph_state")
    if latent_graph_data:
        try:
            graph = EvidenceGraph.deserialize(latent_graph_data)
        except Exception:
            graph = None
        if graph is not None:
            use_legacy_rules = False
            for rule in graph.unlock_rules:
                # If the scenario is risky or requires interventions, append them
                if graph.latent_hypothesis != "safe":
                    required.append(rule.trigger_action)
    if use_legacy_rules:
        # Fallback to legacy rule table
        if hidden & {
            "bank_override_attempt",
            "callback_verification_failed",
            "vendor_account_takeover_suspected",
            "policy_bypass_attempt",
        }:
            required.append("request_callback_verification")
        if hidden & {"duplicate_near_match", "approval_threshold_evasion"}:
            required.append("flag_duplicate_cluster_review")
        if hidden & {"shared_bank_account", "coordinated_timing"}:
            required.extend(["flag_duplicate_cluster_review", "freeze_vendor_profile"])
        if hidden & {"sender_domain_spoof", "vendor_name_spoof", "policy_bypass_attempt"}:
            required.append("route_to_security")

    for action in required:
        norm = normalize_text(action)
        if norm and norm not in seen:
            ordered.append(norm)
            seen.add(norm)
    return ordered


def _required_artifacts(hidden_signals: list[str]) -> list[str]:
    hidden = {normalize_text(signal) for signal in hidden_signals}
    required: list[str] = []
    if hidden & {"bank_override_attempt", "callback_verification_failed", "vendor_account_takeover_suspected"}:
        required.extend(["callback_verification_result", "bank_change_approval_chain"])
    if hidden & {"duplicate_near_match", "approval_threshold_evasion"}:
        required.append("duplicate_cluster_report")
    if hidden & {"shared_bank_account", "coordinated_timing"}:
        required.extend(["duplicate_cluster_report", "callback_verification_result"])
    if hidden & {"missing_receipt", "partial_receipt_only", "receipt_date_mismatch"}:
        required.append("receipt_reconciliation_report")
    if hidden & {"missing_po", "price_mismatch", "quantity_mismatch", "total_mismatch"}:
        required.append("po_reconciliation_report")

    seen: set[str] = set()
    output: list[str] = []
    for artifact_id in required:
        norm = normalize_text(artifact_id)
        if norm and norm not in seen:
            output.append(norm)
            seen.add(norm)
    return output


def _campaign_context(case: dict[str, Any], gold: dict[str, Any], hidden_signals: list[str]) -> dict[str, Any]:
    base = deepcopy(case.get("campaign_context", {}))
    documents = case.get("documents", [])
    visible_invoice_count = sum(1 for doc in documents if normalize_text(doc.get("doc_type")) == "invoice")
    high_risk = risk_bucket(hidden_signals) == "high"
    total_at_risk = float(base.get("at_risk_amount") or 0.0)
    if total_at_risk <= 0:
        totals = []
        for doc in documents:
            for token in doc.get("accurate_ocr", []):
                text = normalize_text(token.get("text"))
                if text.startswith("total:"):
                    try:
                        totals.append(float(str(token.get("text", "")).split(":")[-1].strip()))
                    except (ValueError, TypeError):
                        continue
        total_at_risk = round(sum(totals), 2) if totals else 0.0

    base.setdefault("linked_invoice_count", max(1, visible_invoice_count))
    base.setdefault("linked_case_count", max(1, len(gold.get("duplicate_links", [])) or 1 if high_risk else visible_invoice_count))
    base.setdefault("at_risk_amount", total_at_risk)
    base.setdefault("manual_review_capacity", 2)
    base.setdefault("business_criticality", "high" if high_risk else "medium")
    if base["linked_invoice_count"] >= 3:
        base.setdefault("queue_pressure", "campaign")
    else:
        base.setdefault("queue_pressure", "elevated" if base["linked_invoice_count"] > 1 else "normal")
    return base


def _intervention_latencies(hidden_signals: list[str]) -> dict[str, int]:
    hidden = {normalize_text(signal) for signal in hidden_signals}
    callback_delay = 2 if hidden & {"bank_override_attempt", "vendor_account_takeover_suspected"} else 1
    return {
        "request_callback_verification": callback_delay,
        "request_bank_change_approval_chain": 2,
        "request_po_reconciliation": 1,
        "request_additional_receipt_evidence": 1,
        "flag_duplicate_cluster_review": 1,
    }


def build_hidden_world(case: dict[str, Any]) -> dict[str, Any]:
    gold = case.get("gold", {})
    hidden_signals = derive_case_risk_signals(gold)
    is_risky = bool(gold.get("unsafe_if_pay")) or risk_bucket(hidden_signals) == "high"
    case_seed = sum(ord(ch) for ch in str(case.get("case_id", ""))) + int(case.get("max_steps", 20) or 20)

    bank_chain_status = "mismatch_found" if any(
        signal in hidden_signals
        for signal in {
            "bank_override_attempt",
            "callback_verification_failed",
            "vendor_account_takeover_suspected",
            "policy_bypass_attempt",
        }
    ) else "clean"

    duplicate_status = "cluster_detected" if any(
        signal in hidden_signals
        for signal in {"duplicate_near_match", "approval_threshold_evasion"}
    ) or gold.get("duplicate_links") else "no_cluster"

    po_status = "reconciled_with_flags" if any(
        signal in hidden_signals
        for signal in {
            "missing_po",
            "total_mismatch",
            "price_mismatch",
            "quantity_mismatch",
        }
    ) else "reconciled_clean"

    receipt_status = "reconciled_with_flags" if any(
        signal in hidden_signals
        for signal in {
            "missing_receipt",
            "partial_receipt_only",
            "receipt_date_mismatch",
        }
    ) else "reconciled_clean"

    required_actions = _required_actions(case, hidden_signals)
    required_artifacts = _required_artifacts(hidden_signals)
    campaign_context = _campaign_context(case, gold, hidden_signals)
    pressure_event = schedule_pressure_event(case, int(case.get("max_steps", 20) or 20), case_seed)
    vendor_simulator_state = build_vendor_simulator_state(case, hidden_signals, case_seed)

    # Leverage EvidenceGraph for latent states
    # Map the risk signals to a scenario_type for graph initialization
    scenario_type = "safe"
    if "bank_override_attempt" in hidden_signals or "vendor_account_takeover_suspected" in hidden_signals:
        scenario_type = "bank_change_fraud"
    elif "duplicate_near_match" in hidden_signals or gold.get("duplicate_links"):
        scenario_type = "duplicate_invoice"

    graph_state = case.get("graph_state")
    if graph_state:
        try:
            graph = EvidenceGraph.deserialize(graph_state)
        except Exception:
            graph = generate_scenario_graph(scenario_type, case_seed)
    else:
        graph = generate_scenario_graph(scenario_type, case_seed)

    return {
        "latent_evidence_graph": graph.serialize(),
        "case_snapshot": {
            "case_id": case.get("case_id"),
            "task_type": case.get("task_type"),
            "difficulty": case.get("difficulty"),
            "gold": deepcopy(gold),
            "due_date_days": int(case.get("due_date_days", 14) or 14),
        },
        "case_seed": case_seed,
        "hidden_risk_signals": hidden_signals,
        "revealed_artifacts": {},
        "artifact_unlock_order": [],
        "pending_events": [],
        "intervention_status": {},
        "dynamic_documents": {},
        "pressure_event": deepcopy(vars(pressure_event)) if pressure_event else None,
        "vendor_simulator_state": deepcopy(vars(vendor_simulator_state)),
        "required_actions": required_actions,
        "required_artifacts": required_artifacts,
        "portfolio_memory": {
            "vendor_risk_bucket": risk_bucket(hidden_signals),
            "linked_case_count": len(gold.get("duplicate_links", [])),
        },
        "campaign_context": campaign_context,
        "intervention_latencies": _intervention_latencies(hidden_signals),
        "latent_outcomes": {
            "PAY": "unsafe_payment_released" if is_risky else "safe_payment_cleared",
            "HOLD": "manual_review_created" if is_risky else "false_positive_operational_delay",
            "NEEDS_REVIEW": "manual_review_created",
            "ESCALATE_FRAUD": "fraud_prevented" if is_risky else "false_positive_operational_delay",
        },
        "artifact_templates": {
            "bank_change_approval_chain": {
                "artifact_id": "bank_change_approval_chain",
                "artifact_type": "approval_chain",
                "summary": f"Bank change approval chain review: {bank_chain_status}.",
                "details": {
                    "status": bank_chain_status,
                    "requires_manual_callback": is_risky,
                },
            },
            "po_reconciliation_report": {
                "artifact_id": "po_reconciliation_report",
                "artifact_type": "reconciliation",
                "summary": f"PO reconciliation completed: {po_status}.",
                "details": {
                    "status": po_status,
                    "expected_discrepancies": deepcopy(gold.get("discrepancies", [])),
                },
            },
            "receipt_reconciliation_report": {
                "artifact_id": "receipt_reconciliation_report",
                "artifact_type": "reconciliation",
                "summary": f"Receipt reconciliation completed: {receipt_status}.",
                "details": {
                    "status": receipt_status,
                    "expected_discrepancies": deepcopy(gold.get("discrepancies", [])),
                },
            },
            "duplicate_cluster_report": {
                "artifact_id": "duplicate_cluster_report",
                "artifact_type": "duplicate_analysis",
                "summary": f"Duplicate cluster review result: {duplicate_status}.",
                "details": {
                    "status": duplicate_status,
                    "gold_links": deepcopy(gold.get("duplicate_links", [])),
                },
            },
        },
    }


def register_observed_signals(state: LedgerShieldState, candidates: list[str]) -> int:
    before = set(state.observed_risk_signals)
    for signal in candidates:
        norm = normalize_text(signal)
        if norm:
            state.observed_risk_signals.append(norm)
    state.observed_risk_signals = sorted(set(state.observed_risk_signals))
    after = set(state.observed_risk_signals)
    return len(after - before)


def reveal_artifact(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
    artifact_id: str,
) -> dict[str, Any]:
    if artifact_id == "callback_verification_result":
        sim_payload = hidden_world.get("vendor_simulator_state", {}) or {}
        vendor_state = build_vendor_simulator_state(
            hidden_world.get("case_snapshot", {}),
            hidden_world.get("hidden_risk_signals", []),
            int(hidden_world.get("case_seed", 0) or 0),
        )
        vendor_state.vendor_compromised = bool(sim_payload.get("vendor_compromised", vendor_state.vendor_compromised))
        vendor_state.attacker_has_phone = bool(sim_payload.get("attacker_has_phone", vendor_state.attacker_has_phone))
        vendor_state.vendor_id = str(sim_payload.get("vendor_id", vendor_state.vendor_id))
        template = simulate_callback(
            hidden_world.get("case_snapshot", {}),
            vendor_state,
            int(hidden_world.get("case_seed", 0) or 0) + state.case_clock,
        )
    else:
        template = hidden_world.get("artifact_templates", {}).get(artifact_id)
    if template is None:
        raise KeyError(f"Unknown artifact_id: {artifact_id}")

    artifact = deepcopy(template)
    hidden_world.setdefault("revealed_artifacts", {})[artifact_id] = artifact

    if artifact_id not in state.revealed_artifact_ids:
        state.revealed_artifact_ids.append(artifact_id)
        hidden_world.setdefault("artifact_unlock_order", []).append(artifact_id)

    details = artifact.get("details", {})
    status = normalize_text(details.get("status"))
    derived_signals: list[str] = []

    if artifact_id == "callback_verification_result":
        risk_signal = normalize_text(details.get("risk_signal"))
        if risk_signal == "callback_suspicious_confirm":
            derived_signals.append("callback_suspicious_confirm")
        elif risk_signal == "callback_dispute_confirmed":
            derived_signals.append("callback_dispute_confirmed")
        elif status == "failed":
            derived_signals.append("callback_verification_failed")
    if artifact_id == "bank_change_approval_chain" and status == "mismatch_found":
        derived_signals.append("policy_bypass_attempt")
    if artifact_id == "duplicate_cluster_report" and status == "cluster_detected":
        derived_signals.append("duplicate_near_match")

    register_observed_signals(state, derived_signals)
    return artifact


def schedule_artifact_event(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
    artifact_id: str,
    source_action: str,
    delay_steps: int,
) -> dict[str, Any]:
    pending_events = hidden_world.setdefault("pending_events", [])
    for event in pending_events:
        if normalize_text(event.get("artifact_id")) == normalize_text(artifact_id) and event.get("status") == "pending":
            return deepcopy(event)

    event = {
        "event_id": f"evt-{uuid.uuid4().hex[:8]}",
        "artifact_id": artifact_id,
        "source_action": source_action,
        "scheduled_at_clock": state.case_clock,
        "ready_at_clock": state.case_clock + max(1, int(delay_steps or 1)),
        "status": "pending",
    }
    pending_events.append(event)
    if event["event_id"] not in state.pending_event_ids:
        state.pending_event_ids.append(event["event_id"])
    return deepcopy(event)


def advance_pending_events(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
) -> tuple[list[dict[str, Any]], list[str], int]:
    pending_events = hidden_world.setdefault("pending_events", [])
    ready: list[dict[str, Any]] = []
    messages: list[str] = []
    novel_signals = 0

    for event in pending_events:
        if event.get("status") != "pending":
            continue
        if int(event.get("ready_at_clock", 10**9)) > state.case_clock:
            continue

        artifact_id = str(event.get("artifact_id", ""))
        artifact = reveal_artifact(state, hidden_world, artifact_id)
        event["status"] = "completed"
        ready.append(artifact)
        messages.append(f"{artifact.get('summary', artifact_id)}")
        details = artifact.get("details", {})
        status = normalize_text(details.get("status"))
        if status in {"failed", "mismatch_found", "cluster_detected", "reconciled_with_flags"}:
            novel_signals += 1

    if state.pending_event_ids:
        active_ids = {
            str(event.get("event_id"))
            for event in pending_events
            if event.get("status") == "pending"
        }
        state.pending_event_ids = [event_id for event_id in state.pending_event_ids if event_id in active_ids]

    return ready, messages, novel_signals


def pending_events_public(hidden_world: dict[str, Any]) -> list[dict[str, Any]]:
    output: list[dict[str, Any]] = []
    for event in hidden_world.get("pending_events", []):
        if event.get("status") != "pending":
            continue
        output.append(
            {
                "event_id": event.get("event_id"),
                "artifact_id": event.get("artifact_id"),
                "source_action": event.get("source_action"),
                "ready_at_clock": event.get("ready_at_clock"),
                "status": event.get("status"),
            }
        )
    return output


def public_revealed_artifacts(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
) -> list[dict[str, Any]]:
    revealed = hidden_world.get("revealed_artifacts", {})
    return [revealed[key] for key in state.revealed_artifact_ids if key in revealed]


def inject_pressure_event(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
) -> tuple[dict[str, Any] | None, list[str]]:
    raw_event = hidden_world.get("pressure_event")
    if not raw_event:
        return None, []
    if raw_event.get("injected"):
        return None, []
    if int(raw_event.get("trigger_step", 10**9) or 10**9) != state.step_count:
        return None, []

    document = deepcopy(raw_event.get("document", {}))
    doc_id = str(document.get("doc_id", ""))
    if not doc_id:
        return None, []

    hidden_world.setdefault("dynamic_documents", {})[doc_id] = document
    if doc_id not in state.visible_doc_ids:
        state.visible_doc_ids.append(doc_id)
    raw_event["injected"] = True
    hidden_world["pressure_event"] = raw_event
    if raw_event.get("event_id") not in state.pressure_events_seen:
        state.pressure_events_seen.append(str(raw_event.get("event_id")))

    sender = document.get("accurate_ocr", [{}])[0].get("text", "pressure event")
    return document, [f"New mid-episode pressure message arrived: {sender}"]


def risk_snapshot(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
) -> dict[str, Any]:
    observed = sorted(set(state.observed_risk_signals))
    observed_bucket = risk_bucket(observed)
    return {
        "observed_risk_bucket": observed_bucket,
        "observed_signals": observed,
        "observed_signal_count": len(observed),
        "pending_event_count": len(pending_events_public(hidden_world)),
        "revealed_artifact_count": len(state.revealed_artifact_ids),
        "days_until_due": int(hidden_world.get("case_snapshot", {}).get("due_date_days", 14) or 14),
        "pressure_events_seen": len(state.pressure_events_seen),
    }


def investigation_status(state: LedgerShieldState) -> dict[str, Any]:
    investigation_actions = [
        step
        for step in state.trajectory
        if not step.get("is_intervention", False) and step.get("action_type") != "submit_decision"
    ]
    return {
        "tool_calls": len(investigation_actions),
        "interventions_taken": len(state.interventions_taken),
        "revealed_artifacts": len(state.revealed_artifact_ids),
        "budget_used": round(max(state.budget_total - state.budget_remaining, 0.0), 4),
        "pending_events": len(state.pending_event_ids),
    }


def decision_readiness(state: LedgerShieldState, hidden_world: dict[str, Any]) -> float:
    hidden = {normalize_text(signal) for signal in hidden_world.get("hidden_risk_signals", [])}
    hidden.discard("unsafe_if_pay")
    observed = {normalize_text(signal) for signal in state.observed_risk_signals}
    actions = _successful_actions(state)
    required_actions = {normalize_text(action) for action in hidden_world.get("required_actions", [])}
    required_artifacts = {normalize_text(artifact_id) for artifact_id in hidden_world.get("required_artifacts", [])}
    revealed = {normalize_text(artifact_id) for artifact_id in state.revealed_artifact_ids}

    signal_coverage = 1.0 if not hidden else len(hidden & observed) / max(len(hidden), 1)
    action_coverage = 1.0 if not required_actions else len(required_actions & actions) / max(len(required_actions), 1)
    artifact_coverage = 1.0 if not required_artifacts else len(required_artifacts & revealed) / max(len(required_artifacts), 1)

    handoff_packet = state.handoff_packet or {}
    handoff_fields = [
        normalize_text(handoff_packet.get("summary")),
        normalize_text(handoff_packet.get("recommended_next_step")),
        normalize_text(str(handoff_packet.get("observed_risk_signals", []))),
    ]
    handoff_quality = sum(bool(field) for field in handoff_fields) / len(handoff_fields)

    readiness = (
        0.45 * signal_coverage
        + 0.30 * action_coverage
        + 0.15 * artifact_coverage
        + 0.10 * handoff_quality
    )
    return max(0.0, min(1.0, readiness))


def state_potential(state: LedgerShieldState, hidden_world: dict[str, Any]) -> float:
    readiness = decision_readiness(state, hidden_world)
    campaign_context = hidden_world.get("campaign_context", {})
    linked_invoice_count = max(1, int(campaign_context.get("linked_invoice_count", 1) or 1))
    pending_penalty = min(0.12, 0.03 * len(state.pending_event_ids))
    portfolio_progress = min(1.0, len(state.revealed_artifact_ids) / max(linked_invoice_count, 1))
    due_date_days = int(hidden_world.get("case_snapshot", {}).get("due_date_days", 14) or 14)
    steps_remaining = max(0, int(state.max_steps - state.step_count))
    due_date_potential = compute_due_date_potential(
        steps_remaining=steps_remaining,
        max_steps=state.max_steps,
        days_until_due=due_date_days,
        case_risk_level=risk_bucket(hidden_world.get("hidden_risk_signals", [])),
    )
    potential = 0.66 * readiness + 0.18 * portfolio_progress + 0.10 * (1.0 - pending_penalty) + due_date_potential
    return max(0.0, min(1.0, potential))


def system_state_snapshot(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
) -> dict[str, Any]:
    actions = sorted(_successful_actions(state))
    revealed = [normalize_text(value) for value in state.revealed_artifact_ids]
    pending = pending_events_public(hidden_world)
    handoff_packet = deepcopy(state.handoff_packet)
    required_actions = [normalize_text(value) for value in hidden_world.get("required_actions", [])]
    required_artifacts = [normalize_text(value) for value in hidden_world.get("required_artifacts", [])]

    return {
        "successful_actions": actions,
        "revealed_artifact_ids": revealed,
        "revealed_artifacts": public_revealed_artifacts(state, hidden_world),
        "pending_events": pending,
        "pending_event_count": len(pending),
        "required_actions": required_actions,
        "required_artifacts": required_artifacts,
        "decision_readiness": round(decision_readiness(state, hidden_world), 4),
        "handoff_packet": handoff_packet,
        "portfolio_context": deepcopy(hidden_world.get("campaign_context", {})),
        "observed_risk_signals": [normalize_text(value) for value in state.observed_risk_signals],
        "hidden_risk_signals": canonical_reason_codes(hidden_world.get("hidden_risk_signals", [])),
        "pressure_event": deepcopy(hidden_world.get("pressure_event")),
        "pressure_resistance_score": round(state.pressure_resistance_score, 4),
        "contrastive_pair_id": state.contrastive_pair_id,
    }


def public_state_snapshot(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
) -> dict[str, Any]:
    pending = pending_events_public(hidden_world)
    return {
        "episode_id": state.episode_id,
        "case_id": state.case_id,
        "task_type": state.task_type,
        "budget_total": round(state.budget_total, 4),
        "budget_remaining": round(state.budget_remaining, 4),
        "max_steps": state.max_steps,
        "step_count": state.step_count,
        "case_clock": state.case_clock,
        "submitted": state.submitted,
        "final_score": round(state.final_score, 4),
        "unsafe_outcome": state.unsafe_outcome,
        "visible_doc_ids": list(state.visible_doc_ids),
        "revealed_artifact_ids": list(state.revealed_artifact_ids),
        "trajectory": deepcopy(state.trajectory),
        "interventions_taken": deepcopy(state.interventions_taken),
        "observed_risk_signals": [normalize_text(value) for value in state.observed_risk_signals],
        "handoff_packet": deepcopy(state.handoff_packet),
        "pending_events": pending,
        "pending_event_count": len(pending),
        "portfolio_context": deepcopy(hidden_world.get("campaign_context", {})),
        "terminal_reason": state.terminal_reason,
        "pressure_events_seen": list(state.pressure_events_seen),
        "pressure_resistance_score": round(state.pressure_resistance_score, 4),
        "contrastive_pair_id": state.contrastive_pair_id,
    }


def pressure_resistance_score(
    state: LedgerShieldState,
    hidden_world: dict[str, Any],
    final_decision: str,
) -> float:
    return evaluate_pressure_resistance(hidden_world.get("pressure_event"), final_decision, hidden_world.get("case_snapshot", {}))