"""ROS2 Driver - Real robot control via ROS 2 / rclpy."""

import logging
from typing import Any, Optional

from rosclaw.mcp_drivers.base import BaseDriver, DriverState, TrajectoryCommand

logger = logging.getLogger("rosclaw.mcp_drivers.ros2")


class ROS2Driver(BaseDriver):
    """
    Hardware driver for ROS 2 controlled robots.

    Communicates with the robot via ROS 2 topics and services.
    Requires rclpy and ROS 2 environment to be available.
    """

    def __init__(self, robot_id: str, joint_dof: int = 6, node_name: str = "rosclaw_driver"):
        super().__init__(robot_id, joint_dof)
        self._node_name = node_name
        self._rclpy: Optional[Any] = None
        self._node: Optional[Any] = None
        self._pub_joint_cmd: Optional[Any] = None
        self._sub_joint_state: Optional[Any] = None
        self._latest_joint_state: Optional[dict] = None

    def _do_initialize(self) -> None:
        try:
            import rclpy
            from rclpy.node import Node
            from trajectory_msgs.msg import JointTrajectory
            from sensor_msgs.msg import JointState

            self._rclpy = rclpy
            if not rclpy.ok():
                rclpy.init(args=None)
            self._node = Node(self._node_name)

            self._pub_joint_cmd = self._node.create_publisher(
                JointTrajectory, "/joint_trajectory_controller/joint_trajectory", 10
            )
            self._sub_joint_state = self._node.create_subscription(
                JointState, "/joint_states", self._on_joint_state, 10
            )
            self._driver_state.connected = True
            logger.info("ROS 2 node '%s' initialized", self._node_name)
        except Exception as e:
            logger.warning("rclpy unavailable or init failed (%s), running in mock mode", e)
            self._rclpy = None
            self._node = None
            self._driver_state.connected = True

    def _do_stop(self) -> None:
        if self._node:
            self._node.destroy_node()
            self._node = None
        # NOTE: Do NOT call rclpy.shutdown() here — it's global state.
        # Multiple test instances share the same rclpy context.
        self._driver_state.connected = False

    def _on_joint_state(self, msg: Any) -> None:
        self._latest_joint_state = {
            "positions": list(msg.position),
            "velocities": list(msg.velocity),
            "efforts": list(msg.effort),
        }
        self._driver_state.joint_positions = list(msg.position)[:self.joint_dof]
        self._driver_state.joint_velocities = list(msg.velocity)[:self.joint_dof]
        self._driver_state.joint_torques = list(msg.effort)[:self.joint_dof]

    def get_joint_positions(self) -> list[float]:
        if self._latest_joint_state:
            return self._latest_joint_state["positions"][:self.joint_dof]
        return [0.0] * self.joint_dof

    def get_joint_velocities(self) -> list[float]:
        if self._latest_joint_state:
            return self._latest_joint_state["velocities"][:self.joint_dof]
        return [0.0] * self.joint_dof

    def get_joint_torques(self) -> list[float]:
        if self._latest_joint_state:
            return self._latest_joint_state["efforts"][:self.joint_dof]
        return [0.0] * self.joint_dof

    def move_joints(self, positions: list[float], duration: float = 2.0) -> bool:
        self._ensure_ready("move_joints")
        self._validate_joint_positions(positions)
        self._validate_duration(duration)
        if not self._driver_state.connected:
            return False

        if self._rclpy is None:
            self._driver_state.joint_positions = list(positions)
            return True

        from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint

        traj = JointTrajectory()
        traj.joint_names = [f"joint_{i}" for i in range(self.joint_dof)]
        point = JointTrajectoryPoint()
        point.positions = positions
        point.time_from_start.sec = int(duration)
        traj.points.append(point)
        self._pub_joint_cmd.publish(traj)
        self._driver_state.joint_positions = list(positions)
        return True

    def execute_trajectory(self, trajectory: TrajectoryCommand) -> bool:
        if not self._driver_state.connected:
            return False
        self._ensure_ready("execute_trajectory")
        self._validate_trajectory(trajectory)
        if self._rclpy is None:
            # Mock mode: update state and return success
            if trajectory.waypoints:
                self._driver_state.joint_positions = list(trajectory.waypoints[-1])
            return True
        from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint

        traj = JointTrajectory()
        traj.joint_names = [f"joint_{i}" for i in range(self.joint_dof)]
        for wp, t in zip(trajectory.waypoints, trajectory.times):
            point = JointTrajectoryPoint()
            point.positions = wp
            point.time_from_start.sec = int(t)
            traj.points.append(point)
        self._pub_joint_cmd.publish(traj)
        return True

    def set_gripper(self, position: float, force: float = 0.5) -> bool:
        self._driver_state.gripper_state = position
        return True

    def emergency_stop(self) -> None:
        self._driver_state.error_code = 99
        self._driver_state.error_message = "Emergency stop triggered"
        self.move_joints(self.get_joint_positions(), duration=0.1)

    def get_state(self) -> DriverState:
        return self._driver_state
