"""ONNX policy wrapper with the training-side metadata."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

try:
    import onnxruntime as ort
except Exception:  # pragma: no cover - optional at import time
    ort = None


def _csv_floats(value: str | None) -> list[float]:
    return [float(x) for x in value.split(",") if x.strip()] if value else []


def _csv_strings(value: str | None) -> list[str]:
    return [x.strip() for x in value.split(",") if x.strip()] if value else []


@dataclass
class PolicyMetadata:
    joint_names: list[str]          # Isaac articulation order, with "_joint" suffix
    action_scale: list[float]
    default_joint_pos: list[float]
    body_names: list[str]
    observation_names: list[str]
    raw: dict


class OnnxPolicy:
    """``actions = policy(obs)``; the observation normaliser is baked into the graph."""

    def __init__(self, path: str, device: str = "cpu"):
        if ort is None:
            raise ImportError("onnxruntime is required to run an exported policy (pip install onnxruntime).")
        available = set(ort.get_available_providers())
        if device == "cuda" and "CUDAExecutionProvider" in available:
            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        elif device in ("cpu", "cuda"):
            providers = ["CPUExecutionProvider"]
            if device == "cuda":
                print("[WARN] ONNX Runtime has no CUDA provider; the 182-D policy will run on CPU.")
        else:
            raise ValueError(f"unsupported policy device {device!r}; expected 'cpu' or 'cuda'")
        self.session = ort.InferenceSession(path, providers=providers)
        inputs = [i.name for i in self.session.get_inputs()]
        outputs = [o.name for o in self.session.get_outputs()]
        if inputs != ["obs"] or "actions" not in outputs:
            raise RuntimeError(f"expected an action-only policy with input ['obs'] and output 'actions', got {inputs}/{outputs}")
        self.obs_dim = int(self.session.get_inputs()[0].shape[-1])
        self.act_dim = int([o for o in self.session.get_outputs() if o.name == "actions"][0].shape[-1])
        meta = self.session.get_modelmeta().custom_metadata_map
        self.metadata = PolicyMetadata(
            joint_names=_csv_strings(meta.get("joint_names")),
            action_scale=_csv_floats(meta.get("action_scale")),
            default_joint_pos=_csv_floats(meta.get("default_joint_pos")),
            body_names=_csv_strings(meta.get("body_names")),
            observation_names=_csv_strings(meta.get("observation_names")),
            raw=dict(meta),
        )
        if not self.metadata.joint_names:
            raise RuntimeError("policy.onnx carries no joint_names metadata; export with attach_onnx_metadata")

    def __call__(self, obs: np.ndarray) -> np.ndarray:
        obs = np.asarray(obs, dtype=np.float32).reshape(-1)
        if obs.size != self.obs_dim or not np.isfinite(obs).all():
            raise ValueError(f"policy observation must be finite shape ({self.obs_dim},), got {obs.shape}")
        out = self.session.run(["actions"], {"obs": obs.reshape(1, -1)})[0]
        action = out.reshape(-1).astype(np.float32)
        if action.size != self.act_dim or not np.isfinite(action).all():
            raise RuntimeError(f"policy returned invalid action shape/values: {action.shape}")
        return action


class HoldPosePolicy:
    """Stand-in policy (zero actions = default pose) for pipeline self-tests without ONNX."""

    def __init__(self, num_actions: int, obs_dim: int, joint_names: list[str], default_joint_pos: list[float], action_scale: list[float]):
        self.obs_dim = obs_dim
        self.act_dim = num_actions
        self.metadata = PolicyMetadata(joint_names, action_scale, default_joint_pos, [], [], {})

    def __call__(self, obs: np.ndarray) -> np.ndarray:
        return np.zeros(self.act_dim, dtype=np.float32)