from launch_ros.actions import Node
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument, OpaqueFunction
from launch.substitutions import (
    LaunchConfiguration,
)
from moveit_configs_utils import MoveItConfigsBuilder, moveit_configs_builder


def generate_launch_description():
    declared_arguments = []
    declared_arguments.append(
        DeclareLaunchArgument(
            "robot_description",
            default_value="''",
            description="robot description param",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "robot_description_semantic",
            default_value="''",
            description="robot description semantic param",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "robot_description_kinematics",
            default_value="''",
            description="robot description kinematics param",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "moveit_config_package",
            default_value="rbs_arm",
            description="MoveIt config package with robot SRDF/XACRO files. Usually the argument \
        is not set, it enables use of a custom moveit config.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "moveit_config_file",
            default_value="rbs_arm.srdf.xacro",
            description="MoveIt SRDF/XACRO description file with the robot.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "use_sim_time",
            default_value="true",
            description="Make MoveIt to use simulation time. This is needed for the trajectory planing in simulation.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "tf_prefix",
            default_value="''",
            description="tf_prefix for robot links",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "namespace",
            default_value="",
            description="Namespace for move_group_node",
        )
    )

    return LaunchDescription(
        declared_arguments + [OpaqueFunction(function=launch_setup)]
    )


def launch_setup(context, *args, **kwargs):
    use_sim_time = LaunchConfiguration("use_sim_time")
    robot_description_content = LaunchConfiguration("robot_description").perform(
        context
    )
    robot_description_semantic_content = LaunchConfiguration(
        "robot_description_semantic"
    ).perform(context)
    namespace = LaunchConfiguration("namespace").perform(context)

    robot_description = {"robot_description": robot_description_content}
    robot_description_semantic = {
        "robot_description_semantic": robot_description_semantic_content
    }
    use_sim_time = {"use_sim_time": use_sim_time}

    moveit_config = (
        MoveItConfigsBuilder(robot_name="rbs_arm", package_name="rbs_mill_assist")
        .trajectory_execution(file_path="moveit/moveit_controllers.yaml")
        .joint_limits("moveit/joint_limits.yaml")
        .pilz_cartesian_limits("moveit/pilz_cartesian_limits.yaml")
        .robot_description_kinematics("moveit/kinematics.yaml")
        .robot_description_semantic("srdf/rbs_arm.srdf", mappings={})
        .planning_pipelines(
            pipelines=["ompl", "stomp", "chomp"], default_planning_pipeline="ompl"
        )
        .to_moveit_configs()
    )
    planners_ids = [
        "RRTstar",
        "PRM",
        "BiTRRT",
        "SPARStwo",
    ]
    (
        moveit_config.planning_pipelines["ompl"]
        .setdefault("arm", {})
        # .setdefault("planner_configs", {})
        .update(
            {"planner_configs":planners_ids}
        )
    )

    move_group_node = Node(
        package="moveit_ros_move_group",
        executable="move_group",
        namespace=namespace,
        parameters=[moveit_config.to_dict(), robot_description, use_sim_time],
    )
    planning_scene_publisher = Node(
        package="rbs_mill_assist",
        executable="planning_scene_publisher",
        parameters=[moveit_config.to_dict(), {"use_sim_time": True}],
    )

    rviz_node = Node(
        package="rviz2",
        executable="rviz2",
        name="rviz2",
        output="log",
        # arguments=["-d", rviz_full_config],
        parameters=[
            use_sim_time,
            moveit_config.robot_description,
            moveit_config.robot_description_semantic,
            moveit_config.planning_pipelines,
            moveit_config.planning_scene_monitor,
            moveit_config.robot_description_kinematics,
            moveit_config.joint_limits,
        ],
    )

    nodes_to_start = [
        move_group_node,
        planning_scene_publisher,
        rviz_node,
    ]

    return nodes_to_start