models.py

from __future__ import annotations

from typing import Any, Literal

from openenv.core.env_server.types import Action, Observation, State
from pydantic import BaseModel, Field


ActionName = Literal[
    "allocate_funds",
    "build_reserve",
    "deploy_reserve",
    "approve_batch",
    "reject_transfer",
    "trigger_field_audit",
    "dispatch_repair",
    "dispatch_repair_team",
    "issue_early_warning",
    "escalate_to_oversight",
]


class ActionCommand(BaseModel):
    name: ActionName
    params: dict[str, Any] = Field(default_factory=dict)


class BDOAction(Action):
    predicted_fraud_level: float | None = Field(
        default=None,
        ge=0.0,
        le=1.0,
        description="District-level fraud estimate inferred from the noisy observation.",
    )
    thought_process: str | None = Field(
        default=None,
        description="Free-form reasoning trace explaining why the batch was chosen.",
    )
    actions: list[ActionCommand] = Field(
        default_factory=list,
        description="Ordered monthly action batch executed until the admin-hour budget runs out.",
    )


class TransferAlertModel(BaseModel):
    transfer_id: str
    village: str
    amount: int
    metadata: dict[str, Any] = Field(default_factory=dict)
    oversight_flag: str
    reason: str


class VillageSnapshot(BaseModel):
    village: str
    reported_demand: int
    report_lag_days: int
    demand_confidence: float
    biometric_signal: float
    stability_index: float
    stability_trend: str


class TreasurySnapshot(BaseModel):
    district_budget: int
    projected_demand: int
    buffer_reserve: int
    emergency_unlocked: bool


class MetaSnapshot(BaseModel):
    month: int
    fiscal_year_remaining: int
    global_event: str
    active_shocks: list[str] = Field(default_factory=list)
    hours_remaining: int
    action_costs: dict[str, int] = Field(default_factory=dict)


class RewardBreakdownModel(BaseModel):
    coverage: float = 0.0
    fraud_prevention: float = 0.0
    solvency: float = 0.0
    stability: float = 0.0
    belief_accuracy: float = 0.0
    total: float = 0.0


class VerifierScoresModel(BaseModel):
    syntactic_verifier: float = 0.0
    rationality_verifier: float = 0.0
    outcome_verifier: float = 0.0


class BDOObservation(Observation):
    meta: MetaSnapshot
    treasury: TreasurySnapshot
    nodes: list[VillageSnapshot]
    high_risk_queue: list[TransferAlertModel] = Field(default_factory=list)
    summary: str = ""
    info: dict[str, Any] = Field(
        default_factory=dict,
        description="Step-level diagnostics, reward breakdowns, and execution details.",
    )


class BDOState(State):
    scenario_name: str = "randomized"
    month: int = 1
    shock_month: int | None = None
    seed: int | None = None
    terminated: bool = False
    last_reward: float | None = None