#!/usr/bin/python3

from rclpy.lifecycle import LifecycleNode, LifecycleState, TransitionCallbackReturn
from rclpy.node import Node
from env_manager.scene import Scene
from env_manager.models.configs import (
    CameraData,
    SceneData,
    RobotData,
    ObjectData,
    MeshObjectData,
    TerrainData,
    LightData,
)
import rclpy
from scenario.bindings.gazebo import GazeboSimulator
from env_manager_interfaces.srv import ResetEnv
from rbs_assets_library import get_world_file


DEFAULT_SCENE: SceneData = SceneData(
    robot=RobotData(
        name="rbs_arm",
        with_gripper=True,
        joint_positioins=[0, 0, 0, 0, 0, 0, 0],
        gripper_jont_positions=[0],
    ),
    objects=[MeshObjectData("bishop", position=(0.0, 1.0, 0.3))],
    camera=[CameraData("robot_camera")],
)


class GazeboRuntime(Node):
    def __init__(self) -> None:
        super().__init__(node_name="rbs_gz_runtime")
        self.declare_parameter("robot_description", "")
        self.gazebo = GazeboSimulator(step_size=0.001, rtf=1.0, steps_per_run=1)

        self.gazebo.insert_world_from_sdf(get_world_file("default"))
        if not self.gazebo.initialize():
            raise RuntimeError("Gazebo cannot be initialized")

        self.scene = Scene(
            node=self,
            gazebo=self.gazebo,
            scene=DEFAULT_SCENE,
            robot_urdf_string=self.get_parameter("robot_description")
            .get_parameter_value()
            .string_value,
        )
        self.scene.init_scene()

        self.reset_env_srv = self.create_service(
            ResetEnv, "/env_manager/reset_env", self.reset_env
        )
        self.is_env_reset = False
        self.timer = self.create_timer(self.gazebo.step_size(), self.gazebo_run)

    def gazebo_run(self):
        if self.is_env_reset:
            self.scene.reset_scene()
            self.is_env_reset = False
        self.gazebo.run()

    def reset_env(self, req, res):
        self.is_env_reset = True
        res.ok = self.is_env_reset
        return res


def main(args=None):
    rclpy.init(args=args)
    runtime_node = GazeboRuntime()
    rclpy.spin(runtime_node)
    runtime_node.gazebo.close()
    runtime_node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()