from __future__ import annotations

import importlib
import os
import sys
from pathlib import Path
from typing import Optional

import numpy as np
import pinocchio as pin

from .kinematics import validate_se3


class RealManFkReference:
    """Offline RealMan Algo FK reference; this class never opens a robot connection."""

    def __init__(self, sdk_root: Optional[Path | str] = None) -> None:
        selected_root = sdk_root or os.environ.get("REALMAN_SDK_ROOT")
        if selected_root is not None:
            root = Path(selected_root).expanduser().resolve()
            if not (root / "Robotic_Arm").is_dir():
                raise FileNotFoundError(
                    f"RealMan SDK root must contain Robotic_Arm/: {root}"
                )
            root_text = str(root)
            if root_text not in sys.path:
                sys.path.insert(0, root_text)
        try:
            module = importlib.import_module("Robotic_Arm.rm_robot_interface")
            ctypes_module = importlib.import_module("Robotic_Arm.rm_ctypes_wrap")
        except ImportError as exc:
            raise ImportError(
                "RealMan API2 Python SDK is unavailable; set REALMAN_SDK_ROOT "
                "or pass sdk_root"
            ) from exc

        self._rm_frame_t = module.rm_frame_t
        self._algo = module.Algo(
            module.rm_robot_arm_model_e.RM_MODEL_RM_75_E,
            module.rm_force_type_e.RM_MODEL_RM_B_E,
        )
        self.api_version = str(ctypes_module.rm_api_version())
        self._active_tool_key: tuple[float, ...] | None = None
        self._set_work_frame_identity()
        self._set_tool_frame(None)

    def _set_work_frame_identity(self) -> None:
        frame = self._rm_frame_t(
            frame_name="s1_work",
            pose=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
            payload=0.0,
            x=0.0,
            y=0.0,
            z=0.0,
        )
        self._algo.rm_algo_set_workframe(frame)

    def _set_tool_frame(self, tool: Optional[pin.SE3]) -> None:
        if tool is None:
            pose = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
        else:
            validate_se3(tool, "tool")
            rpy = pin.rpy.matrixToRpy(tool.rotation)
            pose = tuple(float(value) for value in (*tool.translation, *rpy))
        key = tuple(round(value, 12) for value in pose)
        if key == self._active_tool_key:
            return
        frame = self._rm_frame_t(
            frame_name="s1_tool",
            pose=pose,
            payload=0.0,
            x=0.0,
            y=0.0,
            z=0.0,
        )
        self._algo.rm_algo_set_toolframe(frame)
        self._active_tool_key = key

    def forward(self, q_rad: np.ndarray, tool: Optional[pin.SE3] = None) -> pin.SE3:
        q = np.asarray(q_rad, dtype=float)
        if q.shape != (7,) or not np.all(np.isfinite(q)):
            raise ValueError("RealMan FK configuration must be a finite shape-(7,) vector")
        self._set_tool_frame(tool)
        pose = self._algo.rm_algo_forward_kinematics(np.rad2deg(q).tolist(), flag=0)
        if len(pose) != 7 or not np.all(np.isfinite(pose)):
            raise RuntimeError(f"RealMan Algo returned an invalid FK pose: {pose!r}")
        quaternion_values = np.asarray(pose[3:7], dtype=float)
        norm = float(np.linalg.norm(quaternion_values))
        if norm <= 0.0:
            raise RuntimeError("RealMan Algo returned a zero quaternion")
        qw, qx, qy, qz = quaternion_values / norm
        quaternion = pin.Quaternion(qw, qx, qy, qz)
        return pin.SE3(quaternion.matrix(), np.asarray(pose[:3], dtype=float))

