"""Adaptive online budget-escalation policy using LinUCB.

Implements a contextual-bandit controller that decides per-attempt whether
to keep or increase the request budget.  Supports ``off`` / ``shadow`` /
``on`` modes with epsilon-greedy exploration and warm-up uniform sampling.
"""

from __future__ import annotations

import random
import threading
import uuid
from collections.abc import Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

from . import _json as json
from ._policy_utils import (
    coerce_float as _coerce_float,
)
from ._policy_utils import (
    coerce_nonnegative_int as _coerce_nonnegative_int,
)
from ._policy_utils import (
    safe_float_tracked as _safe_float_tracked,
)
from ._policy_utils import (
    utc_now as _utc_now,
)
from .policy_feature_scaling import (
    DEFAULT_POLICY_FEATURE_SCALING_PROFILE,
    POLICY_FEATURE_SCALING_PROFILES,
    normalize_feature_scaling_profile,
    scale_feature_mapping,
)

ACTION_KEEP_BUDGET = "keep_budget"
ACTION_INCREASE_BUDGET = "increase_budget"
BUDGET_POLICY_ACTIONS = (ACTION_KEEP_BUDGET, ACTION_INCREASE_BUDGET)
BUDGET_POLICY_MODES = {"off", "shadow", "on"}
BUDGET_POLICY_ALGOS = {"linucb"}
STATE_SCHEMA_VERSION = 2
EVENT_SCHEMA_VERSION = 3


@dataclass(frozen=True)
class BudgetPolicyConfig:
    """Immutable configuration for the adaptive budget-escalation policy."""

    mode: str = "off"
    algorithm: str = "linucb"
    alpha: float = 0.6
    epsilon: float = 0.03
    l2: float = 1.0
    warmup_decisions: int = 0
    save_every: int = 50
    feature_scaling_profile: str = DEFAULT_POLICY_FEATURE_SCALING_PROFILE
    state_path: Path | None = None
    events_path: Path | None = None
    random_seed: int | None = None

    @classmethod
    def normalize(
        cls,
        *,
        mode: str = "off",
        algorithm: str = "linucb",
        alpha: float = 0.6,
        epsilon: float = 0.03,
        l2: float = 1.0,
        warmup_decisions: int = 0,
        save_every: int = 50,
        feature_scaling_profile: str = DEFAULT_POLICY_FEATURE_SCALING_PROFILE,
        state_path: Path | None = None,
        events_path: Path | None = None,
        random_seed: int | None = None,
    ) -> BudgetPolicyConfig:
        """Validate and clamp raw parameters into a normalised config."""
        mode_norm = str(mode).strip().lower()
        if mode_norm not in BUDGET_POLICY_MODES:
            mode_norm = "off"
        algo_norm = str(algorithm).strip().lower()
        if algo_norm not in BUDGET_POLICY_ALGOS:
            algo_norm = "linucb"
        alpha_norm = max(0.0, _coerce_float(alpha, default=0.6))
        epsilon_norm = min(1.0, max(0.0, _coerce_float(epsilon, default=0.03)))
        l2_norm = max(1e-9, _coerce_float(l2, default=1.0))
        warmup_norm = _coerce_nonnegative_int(warmup_decisions, default=0)
        save_every_norm = max(1, _coerce_nonnegative_int(save_every, default=50))
        feature_scaling_profile_norm = normalize_feature_scaling_profile(feature_scaling_profile)
        if feature_scaling_profile_norm not in POLICY_FEATURE_SCALING_PROFILES:
            feature_scaling_profile_norm = DEFAULT_POLICY_FEATURE_SCALING_PROFILE
        seed_norm: int | None = None
        if random_seed is not None:
            try:
                seed_norm = int(random_seed)
            except (TypeError, ValueError):
                seed_norm = None
        return cls(
            mode=mode_norm,
            algorithm=algo_norm,
            alpha=alpha_norm,
            epsilon=epsilon_norm,
            l2=l2_norm,
            warmup_decisions=warmup_norm,
            save_every=save_every_norm,
            feature_scaling_profile=feature_scaling_profile_norm,
            state_path=Path(state_path) if state_path is not None else None,
            events_path=Path(events_path) if events_path is not None else None,
            random_seed=seed_norm,
        )


@dataclass(frozen=True)
class BudgetPolicyDecision:
    """Snapshot of a single budget-escalation decision by the policy."""

    decision_id: str
    mode: str
    algorithm: str
    decision_index: int
    created_at_utc: str
    selected_action: str
    applied_action: str
    default_action: str
    propensity: float
    decision_strategy: str
    scores: dict[str, float]
    feature_values: dict[str, float]


@dataclass(frozen=True)
class BudgetPolicyEvent:
    """Logged outcome record pairing a budget decision with its reward."""

    run_id: str
    case_id: str
    condition: str
    attempt: int
    note: str
    decision_id: str
    mode: str
    algorithm: str
    decision_index: int
    selected_action: str
    applied_action: str
    behavior_action: str
    default_action: str
    selected_multiplier: float
    applied_multiplier: float
    default_multiplier: float
    default_controls: dict[str, Any]
    selected_controls: dict[str, Any]
    applied_controls: dict[str, Any]
    selected_controls_changed: dict[str, bool]
    applied_controls_changed: dict[str, bool]
    partial_override: bool
    partial_override_errors: dict[str, str]
    tier_index_before: int
    tier_index_after: int
    propensity: float
    behavior_propensity: float
    decision_strategy: str
    feature_scaling_profile: str
    scores: dict[str, float]
    feature_values: dict[str, float]
    reward_scope: str
    reward: float
    reward_success_hint: float
    reward_progress_hint: float
    reward_progress_score: float
    reward_token_penalty: float
    reward_latency_penalty: float
    reward_stagnation_penalty: float
    reward_budget_exceeded_penalty: float
    request_budget_exceeded: bool
    success: bool
    local_llm_call_count: int
    local_tokens_prompt: int
    local_tokens_completion: int
    local_tokens_total: int
    local_wall_sec: float
    override_reason: str
    model_updated: bool
    update_reason: str
    nan_replacement_counts: dict[str, int]
    created_at_utc: str

    def to_json(self) -> dict[str, Any]:
        """Serialize the event to a JSON-compatible dict."""
        return {
            "schema_version": EVENT_SCHEMA_VERSION,
            "run_id": self.run_id,
            "case_id": self.case_id,
            "condition": self.condition,
            "attempt": self.attempt,
            "note": self.note,
            "decision_id": self.decision_id,
            "mode": self.mode,
            "algorithm": self.algorithm,
            "decision_index": self.decision_index,
            "selected_action": self.selected_action,
            "applied_action": self.applied_action,
            "behavior_action": self.behavior_action,
            "default_action": self.default_action,
            "selected_multiplier": self.selected_multiplier,
            "applied_multiplier": self.applied_multiplier,
            "default_multiplier": self.default_multiplier,
            "default_controls": dict(self.default_controls),
            "selected_controls": dict(self.selected_controls),
            "applied_controls": dict(self.applied_controls),
            "selected_controls_changed": {
                key: bool(value) for key, value in self.selected_controls_changed.items()
            },
            "applied_controls_changed": {
                key: bool(value) for key, value in self.applied_controls_changed.items()
            },
            "partial_override": bool(self.partial_override),
            "partial_override_errors": dict(self.partial_override_errors),
            "tier_index_before": self.tier_index_before,
            "tier_index_after": self.tier_index_after,
            "propensity": self.propensity,
            "behavior_propensity": self.behavior_propensity,
            "decision_strategy": self.decision_strategy,
            "feature_scaling_profile": self.feature_scaling_profile,
            "scores": dict(self.scores),
            "feature_values": dict(self.feature_values),
            "reward_scope": self.reward_scope,
            "reward": self.reward,
            "reward_success_hint": self.reward_success_hint,
            "reward_progress_hint": self.reward_progress_hint,
            "reward_progress_score": self.reward_progress_score,
            "reward_token_penalty": self.reward_token_penalty,
            "reward_latency_penalty": self.reward_latency_penalty,
            "reward_stagnation_penalty": self.reward_stagnation_penalty,
            "reward_budget_exceeded_penalty": self.reward_budget_exceeded_penalty,
            "request_budget_exceeded": self.request_budget_exceeded,
            "success": self.success,
            "local_llm_call_count": self.local_llm_call_count,
            "local_tokens_prompt": self.local_tokens_prompt,
            "local_tokens_completion": self.local_tokens_completion,
            "local_tokens_total": self.local_tokens_total,
            "local_wall_sec": self.local_wall_sec,
            "override_reason": self.override_reason,
            "model_updated": self.model_updated,
            "update_reason": self.update_reason,
            "nan_replacement_counts": dict(self.nan_replacement_counts),
            "created_at_utc": self.created_at_utc,
        }


@dataclass
class LinUCBState:
    """Mutable state for a per-action LinUCB model (A matrices and b vectors)."""

    feature_names: list[str] = field(default_factory=list)
    decisions: int = 0
    updates: int = 0
    action_counts: dict[str, int] = field(
        default_factory=lambda: dict.fromkeys(BUDGET_POLICY_ACTIONS, 0)
    )
    A: dict[str, np.ndarray] = field(default_factory=dict)
    b: dict[str, np.ndarray] = field(default_factory=dict)


class AdaptiveBudgetPolicyController:
    """Thread-safe LinUCB contextual-bandit controller for budget escalation.

    Manages exploration/exploitation via epsilon-greedy with UCB scores,
    persists model state to disk, and logs decision-outcome events as JSONL.
    """

    def __init__(self, config: BudgetPolicyConfig) -> None:
        self.config = config
        self._lock = threading.RLock()
        self._rng = random.Random(config.random_seed)
        self._state = LinUCBState()
        self._save_counter = 0
        self._nan_counts: dict[str, int] = {}
        self.load_error: str | None = None
        if self.config.algorithm == "linucb":
            self._load_state_if_present()

    @property
    def mode(self) -> str:
        """Current policy mode (``off``, ``shadow``, or ``on``)."""
        return self.config.mode

    @property
    def algorithm(self) -> str:
        """Algorithm identifier (currently always ``linucb``)."""
        return self.config.algorithm

    @property
    def state(self) -> LinUCBState:
        """Return a deep copy of the internal LinUCB state."""
        with self._lock:
            return LinUCBState(
                feature_names=list(self._state.feature_names),
                decisions=int(self._state.decisions),
                updates=int(self._state.updates),
                action_counts={k: int(v) for k, v in self._state.action_counts.items()},
                A={k: v.copy() for k, v in self._state.A.items()},
                b={k: v.copy() for k, v in self._state.b.items()},
            )

    def get_nan_stats(self) -> dict[str, int]:
        """Return a snapshot of per-feature NaN replacement counts."""
        with self._lock:
            return dict(self._nan_counts)

    def _tracked_safe_float(self, value: Any, *, feature_name: str, default: float = 0.0) -> float:
        """Convert *value* via :func:`safe_float`, tracking NaN replacements."""
        result, replaced = _safe_float_tracked(value, default=default)
        if replaced:
            self._nan_counts[feature_name] = self._nan_counts.get(feature_name, 0) + 1
        return result

    def decide(
        self,
        *,
        feature_values: Mapping[str, float],
        default_action: str = ACTION_KEEP_BUDGET,
    ) -> BudgetPolicyDecision:
        """Select a budget action given the current context features.

        During warm-up or with probability *epsilon* the action is sampled
        uniformly; otherwise the LinUCB greedy arm is chosen.  In ``shadow``
        mode the selected action is logged but the *default_action* is
        applied.
        """
        cleaned_features = {
            str(k): self._tracked_safe_float(v, feature_name=str(k), default=0.0)
            for k, v in dict(feature_values).items()
        }
        scaled_features = scale_feature_mapping(
            cleaned_features,
            profile=self.config.feature_scaling_profile,
        )
        with self._lock:
            self._ensure_feature_state(cleaned_features=scaled_features)
            decision_index = self._state.decisions + 1
            mode = self.config.mode
            default_norm = (
                default_action if default_action in BUDGET_POLICY_ACTIONS else ACTION_KEEP_BUDGET
            )
            if mode == "off":
                selected = default_norm
                applied = default_norm
                propensity = 1.0
                strategy = "policy_off"
                scores = {
                    ACTION_KEEP_BUDGET: 0.0,
                    ACTION_INCREASE_BUDGET: 0.0,
                }
            else:
                x = self._feature_vector(cleaned_features=scaled_features)
                scores = self._linucb_scores(x=x)
                greedy = max(
                    BUDGET_POLICY_ACTIONS,
                    key=lambda action: (scores[action], action == ACTION_KEEP_BUDGET),
                )
                if decision_index <= self.config.warmup_decisions:
                    selected = self._rng.choice(list(BUDGET_POLICY_ACTIONS))
                    strategy = "warmup_uniform"
                else:
                    if self._rng.random() < self.config.epsilon:
                        selected = self._rng.choice(list(BUDGET_POLICY_ACTIONS))
                        strategy = "epsilon_uniform"
                    else:
                        selected = greedy
                        strategy = "linucb_greedy"
                k_actions = float(len(BUDGET_POLICY_ACTIONS))
                base_prob = self.config.epsilon / k_actions
                if decision_index <= self.config.warmup_decisions:
                    propensity = 1.0 / k_actions
                else:
                    propensity = base_prob + (
                        (1.0 - self.config.epsilon) if selected == greedy else 0.0
                    )
                propensity = min(1.0, max(1e-9, propensity))
                applied = selected if mode == "on" else default_norm

            self._state.decisions = decision_index
            self._state.action_counts[selected] = (
                int(self._state.action_counts.get(selected, 0)) + 1
            )
            return BudgetPolicyDecision(
                decision_id=uuid.uuid4().hex,
                mode=mode,
                algorithm=self.config.algorithm,
                decision_index=decision_index,
                created_at_utc=_utc_now(),
                selected_action=selected,
                applied_action=applied,
                default_action=default_norm,
                propensity=float(propensity),
                decision_strategy=strategy,
                scores={key: float(value) for key, value in scores.items()},
                feature_values=scaled_features,
            )

    def record_outcome(
        self,
        *,
        decision: BudgetPolicyDecision,
        run_id: str,
        case_id: str,
        condition: str,
        attempt: int,
        note: str,
        reward_scope: str,
        reward: float,
        reward_success_hint: float,
        reward_progress_hint: float,
        reward_progress_score: float,
        reward_token_penalty: float,
        reward_latency_penalty: float,
        reward_stagnation_penalty: float,
        reward_budget_exceeded_penalty: float,
        request_budget_exceeded: bool,
        success: bool,
        local_llm_call_count: int,
        local_tokens_prompt: int,
        local_tokens_completion: int,
        local_tokens_total: int,
        local_wall_sec: float,
        selected_multiplier: float,
        applied_multiplier: float,
        default_multiplier: float,
        default_controls: Mapping[str, Any],
        selected_controls: Mapping[str, Any],
        applied_controls: Mapping[str, Any],
        selected_controls_changed: Mapping[str, bool],
        applied_controls_changed: Mapping[str, bool],
        partial_override: bool,
        partial_override_errors: Mapping[str, str],
        tier_index_before: int,
        tier_index_after: int,
        override_reason: str,
    ) -> BudgetPolicyEvent:
        """Record an observed reward for a prior budget decision.

        Updates the LinUCB model when mode is ``on``, appends a JSONL
        event record, and triggers periodic state persistence.
        """
        with self._lock:
            model_updated = False
            update_reason = "policy_not_enabled"
            if self.config.mode == "on" and decision.applied_action in BUDGET_POLICY_ACTIONS:
                model_updated = self._update_model(
                    action=decision.applied_action,
                    reward=reward,
                    feature_values=decision.feature_values,
                )
                update_reason = (
                    "on_policy_update" if model_updated else "update_skipped_incomplete_state"
                )
            elif self.config.mode == "shadow":
                update_reason = "shadow_no_update"
            elif self.config.mode == "off":
                update_reason = "off_no_update"

            selected_action = str(decision.selected_action)
            applied_action = str(decision.applied_action)
            behavior_action = applied_action
            behavior_propensity = 1.0
            if selected_action == applied_action:
                behavior_propensity = min(1.0, max(1e-9, float(decision.propensity)))
            event = BudgetPolicyEvent(
                run_id=str(run_id),
                case_id=str(case_id),
                condition=str(condition),
                attempt=max(0, int(attempt)),
                note=str(note),
                decision_id=str(decision.decision_id),
                mode=str(decision.mode),
                algorithm=str(decision.algorithm),
                decision_index=max(0, int(decision.decision_index)),
                selected_action=selected_action,
                applied_action=applied_action,
                behavior_action=behavior_action,
                default_action=str(decision.default_action),
                selected_multiplier=float(selected_multiplier),
                applied_multiplier=float(applied_multiplier),
                default_multiplier=float(default_multiplier),
                default_controls={str(k): v for k, v in default_controls.items()},
                selected_controls={str(k): v for k, v in selected_controls.items()},
                applied_controls={str(k): v for k, v in applied_controls.items()},
                selected_controls_changed={
                    str(k): bool(v) for k, v in selected_controls_changed.items()
                },
                applied_controls_changed={
                    str(k): bool(v) for k, v in applied_controls_changed.items()
                },
                partial_override=bool(partial_override),
                partial_override_errors={
                    str(k): str(v) for k, v in partial_override_errors.items()
                },
                tier_index_before=max(0, int(tier_index_before)),
                tier_index_after=max(0, int(tier_index_after)),
                propensity=min(1.0, max(1e-9, float(decision.propensity))),
                behavior_propensity=behavior_propensity,
                decision_strategy=str(decision.decision_strategy),
                feature_scaling_profile=str(self.config.feature_scaling_profile),
                scores={k: float(v) for k, v in decision.scores.items()},
                feature_values={k: float(v) for k, v in decision.feature_values.items()},
                reward_scope=str(reward_scope),
                reward=float(reward),
                reward_success_hint=float(reward_success_hint),
                reward_progress_hint=float(reward_progress_hint),
                reward_progress_score=float(reward_progress_score),
                reward_token_penalty=float(reward_token_penalty),
                reward_latency_penalty=float(reward_latency_penalty),
                reward_stagnation_penalty=float(reward_stagnation_penalty),
                reward_budget_exceeded_penalty=float(reward_budget_exceeded_penalty),
                request_budget_exceeded=bool(request_budget_exceeded),
                success=bool(success),
                local_llm_call_count=max(0, int(local_llm_call_count)),
                local_tokens_prompt=max(0, int(local_tokens_prompt)),
                local_tokens_completion=max(0, int(local_tokens_completion)),
                local_tokens_total=max(0, int(local_tokens_total)),
                local_wall_sec=max(0.0, float(local_wall_sec)),
                override_reason=str(override_reason),
                model_updated=bool(model_updated),
                update_reason=update_reason,
                nan_replacement_counts=dict(self._nan_counts),
                created_at_utc=_utc_now(),
            )
            self._append_event(event=event)
            if model_updated:
                self._save_counter += 1
                if self._save_counter >= max(1, int(self.config.save_every)):
                    self._save_counter = 0
                    self._save_state()
            return event

    def close(self) -> None:
        """Flush the current model state to disk."""
        with self._lock:
            self._save_state()

    def _ensure_feature_state(self, *, cleaned_features: Mapping[str, float]) -> None:
        """Initialise or expand A/b matrices when new features appear."""
        if self.config.algorithm != "linucb":
            return
        if not self._state.feature_names:
            self._state.feature_names = sorted(str(k) for k in cleaned_features)
            dim = len(self._state.feature_names)
            if dim <= 0:
                self._state.feature_names = ["bias"]
                dim = 1
            self._state.A = {
                action: (np.eye(dim, dtype=float) * float(self.config.l2))
                for action in BUDGET_POLICY_ACTIONS
            }
            self._state.b = {
                action: np.zeros((dim,), dtype=float) for action in BUDGET_POLICY_ACTIONS
            }
            return
        expanded_feature_names = sorted(
            set(self._state.feature_names) | {str(key) for key in cleaned_features}
        )
        if expanded_feature_names != self._state.feature_names:
            previous_feature_names = list(self._state.feature_names)
            previous_index = {name: idx for idx, name in enumerate(previous_feature_names)}
            previous_dim = len(previous_feature_names)
            new_dim = len(expanded_feature_names)
            self._state.feature_names = list(expanded_feature_names)
            for action in BUDGET_POLICY_ACTIONS:
                previous_A = self._state.A.get(action)
                previous_b = self._state.b.get(action)
                next_A = np.eye(new_dim, dtype=float) * float(self.config.l2)
                next_b = np.zeros((new_dim,), dtype=float)
                if (
                    isinstance(previous_A, np.ndarray)
                    and isinstance(previous_b, np.ndarray)
                    and previous_A.shape == (previous_dim, previous_dim)
                    and previous_b.shape == (previous_dim,)
                ):
                    for new_row, feature_row in enumerate(expanded_feature_names):
                        old_row = previous_index.get(feature_row)
                        if old_row is None:
                            continue
                        next_b[new_row] = previous_b[old_row]
                        for new_col, feature_col in enumerate(expanded_feature_names):
                            old_col = previous_index.get(feature_col)
                            if old_col is None:
                                continue
                            next_A[new_row, new_col] = previous_A[old_row, old_col]
                self._state.A[action] = next_A
                self._state.b[action] = next_b
        dim = len(self._state.feature_names)
        for action in BUDGET_POLICY_ACTIONS:
            A = self._state.A.get(action)
            b = self._state.b.get(action)
            if A is None or A.shape != (dim, dim):
                self._state.A[action] = np.eye(dim, dtype=float) * float(self.config.l2)
            if b is None or b.shape != (dim,):
                self._state.b[action] = np.zeros((dim,), dtype=float)

    def _feature_vector(self, *, cleaned_features: Mapping[str, float]) -> np.ndarray:
        """Build a numpy feature vector aligned with ``feature_names``."""
        feature_names = self._state.feature_names
        if not feature_names:
            feature_names = ["bias"]
        values = np.array(
            [
                self._tracked_safe_float(
                    cleaned_features.get(name, 0.0),
                    feature_name=name,
                    default=0.0,
                )
                for name in feature_names
            ],
            dtype=float,
        )
        return values

    def _linucb_scores(self, *, x: np.ndarray) -> dict[str, float]:
        """Compute UCB scores (exploit + alpha * explore) for every action."""
        scores: dict[str, float] = {}
        if self.config.algorithm != "linucb":
            return dict.fromkeys(BUDGET_POLICY_ACTIONS, 0.0)
        for action in BUDGET_POLICY_ACTIONS:
            A = self._state.A.get(action)
            b = self._state.b.get(action)
            if A is None or b is None:
                scores[action] = 0.0
                continue
            try:
                A_inv = np.linalg.inv(A)
            except np.linalg.LinAlgError:
                A = A + (np.eye(A.shape[0], dtype=float) * float(self.config.l2))
                A_inv = np.linalg.inv(A)
                self._state.A[action] = A
            theta = A_inv @ b
            exploit = float(theta.T @ x)
            explore = float(np.sqrt(max(0.0, float(x.T @ A_inv @ x))))
            scores[action] = exploit + (float(self.config.alpha) * explore)
        return scores

    def _update_model(
        self,
        *,
        action: str,
        reward: float,
        feature_values: Mapping[str, float],
    ) -> bool:
        """Apply a rank-1 LinUCB update: A += x x^T, b += reward * x."""
        if self.config.algorithm != "linucb":
            return False
        if action not in BUDGET_POLICY_ACTIONS:
            return False
        if not self._state.feature_names:
            return False
        x = self._feature_vector(cleaned_features=feature_values)
        A = self._state.A.get(action)
        b = self._state.b.get(action)
        if A is None or b is None:
            return False
        A += np.outer(x, x)
        b += float(reward) * x
        self._state.A[action] = A
        self._state.b[action] = b
        self._state.updates = int(self._state.updates) + 1
        return True

    def _state_payload(self) -> dict[str, Any]:
        """Build a JSON-serialisable snapshot of the full model state."""
        return {
            "schema_version": STATE_SCHEMA_VERSION,
            "mode": self.config.mode,
            "algorithm": self.config.algorithm,
            "alpha": self.config.alpha,
            "epsilon": self.config.epsilon,
            "l2": self.config.l2,
            "warmup_decisions": self.config.warmup_decisions,
            "save_every": self.config.save_every,
            "feature_scaling_profile": self.config.feature_scaling_profile,
            "feature_names": list(self._state.feature_names),
            "decisions": int(self._state.decisions),
            "updates": int(self._state.updates),
            "action_counts": {
                action: int(self._state.action_counts.get(action, 0))
                for action in BUDGET_POLICY_ACTIONS
            },
            "A": {
                action: self._state.A.get(action, np.empty((0, 0), dtype=float)).tolist()
                for action in BUDGET_POLICY_ACTIONS
            },
            "b": {
                action: self._state.b.get(action, np.empty((0,), dtype=float)).tolist()
                for action in BUDGET_POLICY_ACTIONS
            },
            "saved_at_utc": _utc_now(),
        }

    def _save_state(self) -> None:
        """Atomically write model state to disk via a temp-file rename."""
        path = self.config.state_path
        if path is None:
            return
        payload = self._state_payload()
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp_path = path.with_suffix(path.suffix + ".tmp")
        tmp_path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
        tmp_path.replace(path)

    def _load_state_if_present(self) -> None:
        """Load persisted LinUCB state from disk, if it exists."""
        path = self.config.state_path
        if path is None or not path.exists():
            return
        try:
            raw = json.loads(path.read_text(encoding="utf-8"))
            if not isinstance(raw, dict):
                raise ValueError("state payload is not an object")
            if int(raw.get("schema_version", 0)) != STATE_SCHEMA_VERSION:
                raise ValueError("unsupported schema version")
            algorithm = str(raw.get("algorithm") or "")
            if algorithm and algorithm != self.config.algorithm:
                raise ValueError("algorithm mismatch")
            state_scaling_profile = normalize_feature_scaling_profile(
                str(raw.get("feature_scaling_profile") or DEFAULT_POLICY_FEATURE_SCALING_PROFILE)
            )
            config_scaling_profile = normalize_feature_scaling_profile(
                self.config.feature_scaling_profile
            )
            if state_scaling_profile != config_scaling_profile:
                raise ValueError("feature_scaling_profile mismatch")
            feature_names_raw = raw.get("feature_names")
            if not isinstance(feature_names_raw, list) or not all(
                isinstance(item, str) for item in feature_names_raw
            ):
                raise ValueError("missing or invalid feature_names")
            feature_names = [str(item) for item in feature_names_raw]
            dim = len(feature_names)
            decisions = _coerce_nonnegative_int(raw.get("decisions"), default=0)
            updates = _coerce_nonnegative_int(raw.get("updates"), default=0)
            action_counts_raw = raw.get("action_counts")
            action_counts = dict.fromkeys(BUDGET_POLICY_ACTIONS, 0)
            if isinstance(action_counts_raw, dict):
                for action in BUDGET_POLICY_ACTIONS:
                    action_counts[action] = _coerce_nonnegative_int(
                        action_counts_raw.get(action), default=0
                    )
            if dim <= 0:
                self._state = LinUCBState(
                    feature_names=[],
                    decisions=decisions,
                    updates=updates,
                    action_counts=action_counts,
                    A={},
                    b={},
                )
                self.load_error = None
                return
            A_raw = raw.get("A")
            b_raw = raw.get("b")
            if not isinstance(A_raw, dict) or not isinstance(b_raw, dict):
                raise ValueError("missing action matrices")
            A_state: dict[str, np.ndarray] = {}
            b_state: dict[str, np.ndarray] = {}
            for action in BUDGET_POLICY_ACTIONS:
                A_matrix = np.array(A_raw.get(action), dtype=float)
                b_vector = np.array(b_raw.get(action), dtype=float)
                if A_matrix.shape != (dim, dim):
                    raise ValueError(f"invalid A shape for action {action}")
                if b_vector.shape != (dim,):
                    raise ValueError(f"invalid b shape for action {action}")
                A_state[action] = A_matrix
                b_state[action] = b_vector
            self._state = LinUCBState(
                feature_names=feature_names,
                decisions=decisions,
                updates=updates,
                action_counts=action_counts,
                A=A_state,
                b=b_state,
            )
            self.load_error = None
        except Exception as exc:
            self._state = LinUCBState()
            self.load_error = f"{type(exc).__name__}: {exc}"

    def _append_event(self, *, event: BudgetPolicyEvent) -> None:
        """Append one JSONL event line to the events log file."""
        path = self.config.events_path
        if path is None:
            return
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(event.to_json(), ensure_ascii=False))
            handle.write("\n")
