import os
import subprocess

import rospy
from gazebo_msgs.srv import SpawnModel
from geometry_msgs.msg import Point, Pose, Quaternion
from tf.transformations import quaternion_from_euler

from gazebo_ctrl import GazeboROSController


class ModelGenerator:
    def __init__(self):
        self.controller = GazeboROSController()
        ros_package_path = os.environ.get('ROS_PACKAGE_PATH', '')
        self.base_dir = os.path.expanduser('~')
        self.base_path = os.path.join(self.base_dir, 'gazebo_world_manager')
        if self.base_path not in ros_package_path.split(':'):
            os.environ['ROS_PACKAGE_PATH'] = f"{self.base_path}:{ros_package_path}" if ros_package_path else self.base_path

    def generate_table(self, name, x=0, y=0, height=1, roll=0, pitch=0, yaw=0):
        xacro_file = os.path.join(self.base_path, "scene/grasp_box/models/table.xacro")
        urdf_content = subprocess.check_output([
            'rosrun', 'xacro', 'xacro', xacro_file,
            f'x:={x}', f'y:={y}', f'height:={height}'
        ]).decode('utf-8')
        pose = Pose(
            position=Point(x=x, y=y, z=height),
            orientation=Quaternion(*quaternion_from_euler(roll, pitch, yaw))
        )
        return self.controller.spawn_urdf_model(name, urdf_content, pose)

    def generate_apriltag(self, name, tag_id=0, x=0, y=0, z=0, roll=0, pitch=0, yaw=0):
        """生成AprilTag标记模型"""
        xacro_file = os.path.join(self.base_path, "scene/grasp_box/models/fixed_tag.xacro")
        urdf_content = subprocess.check_output([
            'rosrun', 'xacro', 'xacro', xacro_file, f'tag_id:={tag_id}'
        ]).decode('utf-8')
        pose = Pose(
            position=Point(x=x, y=y, z=z),
            orientation=Quaternion(
                *quaternion_from_euler(roll, pitch, yaw))
        )
        return self.controller.spawn_urdf_model(name, urdf_content, pose)

    def generate_box_with_apriltag(self, name, l=0.25, w=0.2, h=0.2, mass=1, tag_id=0, x=0, y=0, z=0, roll=0, pitch=0, yaw=0):
        """生成带有AprilTag的盒子模型"""
        xacro_file = os.path.join(self.base_path, "scene/grasp_box/models/box.xacro")
        urdf_content = subprocess.check_output([
            'rosrun', 'xacro', 'xacro', xacro_file, f'box_length:={l}', f'box_width:={w}', f'box_height:={h}', f'box_mass:={mass}', f'tag_id:={tag_id}'
        ]).decode('utf-8')
        pose = Pose(
            position=Point(x=x, y=y, z=z),
            orientation=Quaternion(
                *quaternion_from_euler(roll, pitch, yaw))
        )
        return self.controller.spawn_urdf_model(name, urdf_content, pose)


if __name__ == '__main__':
    generator = ModelGenerator()
    # generator.generate_apriltag("tg0", 0, 2, 0, 1.1, 0, 1.57, 0)
    generator.generate_table("tb0", x=2, y=1, height=0.8)
    # generator.generate_box_with_apriltag("bx0", x=2, z=1)