from __future__ import annotations

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

import numpy as np

from . import _json as json

ACTION_SKIP_HELPER = "skip_helper"
ACTION_INVOKE_HELPER = "invoke_helper"
ONLINE_POLICY_ACTIONS = (ACTION_SKIP_HELPER, ACTION_INVOKE_HELPER)
ONLINE_POLICY_MODES = {"off", "shadow", "on"}
ONLINE_POLICY_ALGOS = {"linucb"}
STATE_SCHEMA_VERSION = 1


def _utc_now() -> str:
    return (
        datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
    )


def _coerce_float(value: Any, *, default: float) -> float:
    try:
        return float(value)
    except (TypeError, ValueError):
        return float(default)


def _coerce_nonnegative_int(value: Any, *, default: int) -> int:
    try:
        out = int(value)
    except (TypeError, ValueError):
        return int(default)
    return max(0, out)


def _safe_float(value: Any, *, default: float = 0.0) -> float:
    try:
        out = float(value)
    except (TypeError, ValueError):
        return float(default)
    if np.isnan(out) or np.isinf(out):
        return float(default)
    return out


@dataclass(frozen=True)
class OnlinePolicyConfig:
    mode: str = "off"
    algorithm: str = "linucb"
    alpha: float = 0.8
    epsilon: float = 0.05
    l2: float = 1.0
    warmup_decisions: int = 200
    save_every: int = 50
    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.8,
        epsilon: float = 0.05,
        l2: float = 1.0,
        warmup_decisions: int = 200,
        save_every: int = 50,
        state_path: Path | None = None,
        events_path: Path | None = None,
        random_seed: int | None = None,
    ) -> OnlinePolicyConfig:
        mode_norm = str(mode).strip().lower()
        if mode_norm not in ONLINE_POLICY_MODES:
            mode_norm = "off"
        algo_norm = str(algorithm).strip().lower()
        if algo_norm not in ONLINE_POLICY_ALGOS:
            algo_norm = "linucb"
        alpha_norm = max(0.0, _coerce_float(alpha, default=0.8))
        epsilon_norm = min(1.0, max(0.0, _coerce_float(epsilon, default=0.05)))
        l2_norm = max(1e-9, _coerce_float(l2, default=1.0))
        warmup_norm = _coerce_nonnegative_int(warmup_decisions, default=200)
        save_every_norm = max(1, _coerce_nonnegative_int(save_every, default=50))
        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,
            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 PolicyDecision:
    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 PolicyEvent:
    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
    default_action: str
    propensity: float
    decision_strategy: str
    scores: dict[str, float]
    feature_values: dict[str, float]
    reward: float
    reward_success: float
    reward_token_penalty: float
    reward_latency_penalty: float
    success: bool
    tokens_total: int
    attempt_wall_sec: float
    override_reason: str
    model_updated: bool
    update_reason: str
    created_at_utc: str

    def to_json(self) -> dict[str, Any]:
        return {
            "schema_version": STATE_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,
            "default_action": self.default_action,
            "propensity": self.propensity,
            "decision_strategy": self.decision_strategy,
            "scores": dict(self.scores),
            "feature_values": dict(self.feature_values),
            "reward": self.reward,
            "reward_success": self.reward_success,
            "reward_token_penalty": self.reward_token_penalty,
            "reward_latency_penalty": self.reward_latency_penalty,
            "success": self.success,
            "tokens_total": self.tokens_total,
            "attempt_wall_sec": self.attempt_wall_sec,
            "override_reason": self.override_reason,
            "model_updated": self.model_updated,
            "update_reason": self.update_reason,
            "created_at_utc": self.created_at_utc,
        }


@dataclass
class LinUCBState:
    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(ONLINE_POLICY_ACTIONS, 0)
    )
    A: dict[str, np.ndarray] = field(default_factory=dict)
    b: dict[str, np.ndarray] = field(default_factory=dict)


class AdaptiveOnlinePolicyController:
    def __init__(self, config: OnlinePolicyConfig) -> None:
        self.config = config
        self._lock = threading.RLock()
        self._rng = random.Random(config.random_seed)
        self._state = LinUCBState()
        self._save_counter = 0
        self.load_error: str | None = None
        if self.config.algorithm == "linucb":
            self._load_state_if_present()

    @property
    def mode(self) -> str:
        return self.config.mode

    @property
    def algorithm(self) -> str:
        return self.config.algorithm

    @property
    def state(self) -> LinUCBState:
        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 decide(
        self,
        *,
        feature_values: Mapping[str, float],
        default_action: str = ACTION_INVOKE_HELPER,
    ) -> PolicyDecision:
        cleaned_features = {
            str(k): _safe_float(v, default=0.0) for k, v in dict(feature_values).items()
        }
        with self._lock:
            self._ensure_feature_state(cleaned_features=cleaned_features)
            decision_index = self._state.decisions + 1
            mode = self.config.mode
            default_norm = (
                default_action if default_action in ONLINE_POLICY_ACTIONS else ACTION_INVOKE_HELPER
            )
            if mode == "off":
                selected = default_norm
                applied = default_norm
                propensity = 1.0
                strategy = "policy_off"
                scores = {
                    ACTION_SKIP_HELPER: 0.0,
                    ACTION_INVOKE_HELPER: 0.0,
                }
            else:
                x = self._feature_vector(cleaned_features=cleaned_features)
                scores = self._linucb_scores(x=x)
                greedy = max(
                    ONLINE_POLICY_ACTIONS,
                    key=lambda action: (scores[action], action == ACTION_INVOKE_HELPER),
                )
                if decision_index <= self.config.warmup_decisions:
                    selected = self._rng.choice(list(ONLINE_POLICY_ACTIONS))
                    strategy = "warmup_uniform"
                else:
                    if self._rng.random() < self.config.epsilon:
                        selected = self._rng.choice(list(ONLINE_POLICY_ACTIONS))
                        strategy = "epsilon_uniform"
                    else:
                        selected = greedy
                        strategy = "linucb_greedy"
                k_actions = float(len(ONLINE_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 PolicyDecision(
                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=cleaned_features,
            )

    def record_outcome(
        self,
        *,
        decision: PolicyDecision,
        run_id: str,
        case_id: str,
        condition: str,
        attempt: int,
        note: str,
        reward: float,
        reward_success: float,
        reward_token_penalty: float,
        reward_latency_penalty: float,
        success: bool,
        tokens_total: int,
        attempt_wall_sec: float,
        override_reason: str,
    ) -> PolicyEvent:
        with self._lock:
            model_updated = False
            update_reason = "policy_not_enabled"
            if self.config.mode == "on" and decision.applied_action in ONLINE_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"

            event = PolicyEvent(
                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=str(decision.selected_action),
                applied_action=str(decision.applied_action),
                default_action=str(decision.default_action),
                propensity=min(1.0, max(1e-9, float(decision.propensity))),
                decision_strategy=str(decision.decision_strategy),
                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=float(reward),
                reward_success=float(reward_success),
                reward_token_penalty=float(reward_token_penalty),
                reward_latency_penalty=float(reward_latency_penalty),
                success=bool(success),
                tokens_total=max(0, int(tokens_total)),
                attempt_wall_sec=max(0.0, float(attempt_wall_sec)),
                override_reason=str(override_reason),
                model_updated=bool(model_updated),
                update_reason=update_reason,
                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:
        with self._lock:
            self._save_state()

    def _ensure_feature_state(self, *, cleaned_features: Mapping[str, float]) -> None:
        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 ONLINE_POLICY_ACTIONS
            }
            self._state.b = {
                action: np.zeros((dim,), dtype=float) for action in ONLINE_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 ONLINE_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 ONLINE_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:
        feature_names = self._state.feature_names
        if not feature_names:
            feature_names = ["bias"]
        values = np.array(
            [_safe_float(cleaned_features.get(name, 0.0), default=0.0) for name in feature_names],
            dtype=float,
        )
        return values

    def _linucb_scores(self, *, x: np.ndarray) -> dict[str, float]:
        scores: dict[str, float] = {}
        if self.config.algorithm != "linucb":
            return dict.fromkeys(ONLINE_POLICY_ACTIONS, 0.0)
        for action in ONLINE_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:
        if self.config.algorithm != "linucb":
            return False
        if action not in ONLINE_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]:
        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_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 ONLINE_POLICY_ACTIONS
            },
            "A": {
                action: self._state.A.get(action, np.empty((0, 0), dtype=float)).tolist()
                for action in ONLINE_POLICY_ACTIONS
            },
            "b": {
                action: self._state.b.get(action, np.empty((0,), dtype=float)).tolist()
                for action in ONLINE_POLICY_ACTIONS
            },
            "saved_at_utc": _utc_now(),
        }

    def _save_state(self) -> None:
        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:
        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")
            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)
            if dim <= 0:
                raise ValueError("empty feature_names")
            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 matrices")
            A: dict[str, np.ndarray] = {}
            b: dict[str, np.ndarray] = {}
            for action in ONLINE_POLICY_ACTIONS:
                matrix_raw = A_raw.get(action)
                vector_raw = b_raw.get(action)
                matrix = np.array(matrix_raw, dtype=float)
                vector = np.array(vector_raw, dtype=float)
                if matrix.shape != (dim, dim):
                    raise ValueError(f"invalid matrix shape for {action}: {matrix.shape}")
                if vector.shape != (dim,):
                    raise ValueError(f"invalid vector shape for {action}: {vector.shape}")
                A[action] = matrix
                b[action] = vector
            action_counts_raw = raw.get("action_counts")
            action_counts = dict.fromkeys(ONLINE_POLICY_ACTIONS, 0)
            if isinstance(action_counts_raw, dict):
                for action in ONLINE_POLICY_ACTIONS:
                    action_counts[action] = _coerce_nonnegative_int(
                        action_counts_raw.get(action), default=0
                    )
            self._state = LinUCBState(
                feature_names=feature_names,
                decisions=_coerce_nonnegative_int(raw.get("decisions"), default=0),
                updates=_coerce_nonnegative_int(raw.get("updates"), default=0),
                action_counts=action_counts,
                A=A,
                b=b,
            )
            self.load_error = None
        except Exception as exc:  # pragma: no cover - defensive load hardening
            self.load_error = f"{type(exc).__name__}: {exc}"
            self._state = LinUCBState()

    def _append_event(self, *, event: PolicyEvent) -> None:
        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) + "\n")
