from launch import LaunchContext, LaunchDescription
from launch.actions import (
    DeclareLaunchArgument,
    IncludeLaunchDescription,
    RegisterEventHandler,
    OpaqueFunction
)
from launch.event_handlers import OnProcessExit
from launch.conditions import IfCondition
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import Command, FindExecutable, LaunchConfiguration, PathJoinSubstitution
from launch_ros.actions import Node
from launch_ros.substitutions import FindPackageShare
import xacro
import os
from ament_index_python.packages import get_package_share_directory

from rbs_arm import RbsBuilder

def launch_setup(context, *args, **kwargs):
    # Initialize Arguments
    robot_type = LaunchConfiguration("robot_type")
    # General arguments
    with_gripper_condition = LaunchConfiguration("with_gripper")
    controllers_file = LaunchConfiguration("controllers_file")
    cartesian_controllers = LaunchConfiguration("cartesian_controllers")
    description_package = LaunchConfiguration("description_package")
    description_file = LaunchConfiguration("description_file")
    robot_name = LaunchConfiguration("robot_name")
    start_joint_controller = LaunchConfiguration("start_joint_controller")
    initial_joint_controller = LaunchConfiguration("initial_joint_controller")
    launch_moveit = LaunchConfiguration("launch_moveit")
    launch_task_planner = LaunchConfiguration("launch_task_planner")
    launch_perception = LaunchConfiguration("launch_perception")
    moveit_config_package = LaunchConfiguration("moveit_config_package")
    moveit_config_file = LaunchConfiguration("moveit_config_file")
    use_sim_time = LaunchConfiguration("use_sim_time")
    hardware = LaunchConfiguration("hardware")
    launch_controllers = LaunchConfiguration("launch_controllers")
    gripper_name = LaunchConfiguration("gripper_name")
    x_pos = LaunchConfiguration("x")
    y_pos = LaunchConfiguration("y")
    z_pos = LaunchConfiguration("z")
    roll = LaunchConfiguration("roll")
    pitch = LaunchConfiguration("pitch")
    yaw = LaunchConfiguration("yaw")
    namespace = LaunchConfiguration("namespace")
    robot_name = robot_name.perform(context)
    namespace = namespace.perform(context)
    robot_type = robot_type.perform(context)
    description_package = description_package.perform(context)
    description_file = description_file.perform(context)
    controllers_file = controllers_file.perform(context)

    # remappings = [("/tf", "tf"), ("/tf_static", "tf_static")]
    
    xacro_file = os.path.join(get_package_share_directory(description_package), "urdf", description_file)
    robot_description_doc = xacro.process_file(
        xacro_file,
        mappings={
            "gripper_name": gripper_name.perform(context),
            "hardware": hardware.perform(context),
            "simulation_controllers": controllers_file,
            "namespace": namespace,
            "x": x_pos.perform(context),
            "y": y_pos.perform(context),
            "z": z_pos.perform(context),
            "roll": roll.perform(context),
            "pitch": pitch.perform(context),
            "yaw": yaw.perform(context)
            #TODO: add rotation and add probably via dict
        }
    )

    # robot = RbsBuilder(6, "arm0", "world", "rbs_gripper")
    # robot.base()
    # robot.gripper()
    # robot.ros2_control()
    # robot.moveit()
    robot_description_content = robot_description_doc.toprettyxml(indent="  ")

    robot_description = {"robot_description": robot_description_content}
    
    robot_description_semantic_content = Command(
        [
            PathJoinSubstitution([FindExecutable(name="xacro")]),
            " ",
            PathJoinSubstitution(
                [FindPackageShare(moveit_config_package), "config/moveit", "rbs_arm.srdf.xacro"]
            ),
            " ",
            "name:=",robot_name," ",
            "with_gripper:=",with_gripper_condition, " ",
            "gripper_name:=", gripper_name, " ",
        ]
    )
 
    robot_description_semantic = {"robot_description_semantic": robot_description_semantic_content}
    robot_description_kinematics = PathJoinSubstitution(
        [FindPackageShare(moveit_config_package), "config", "kinematics.yaml"]
    )
    robot_state_publisher = Node(
        package="robot_state_publisher",
        executable="robot_state_publisher",
        namespace=namespace,
        output="both",
        # remappings=remappings,
        parameters=[{"use_sim_time": use_sim_time}, robot_description],
    )

    rviz_config_file = PathJoinSubstitution(
        [FindPackageShare("rbs_bringup"), "config", "rbs.rviz"]
    )

    rviz = Node(
        package="rviz2",
        executable="rviz2",
        name="rviz2",
        output="log",
        namespace=namespace,
        arguments=["-d", rviz_config_file],
        parameters=[
            robot_description,
            robot_description_semantic,
            robot_description_kinematics,
            {'use_sim_time': use_sim_time}
        ]
    )

    control = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare(description_package),
                'launch',
                'control.launch.py'
            ])
        ]),
        launch_arguments={
            'with_gripper': with_gripper_condition,
            'robot_description': robot_description_content,
            'start_joint_controller': start_joint_controller,
            'initial_joint_controller': initial_joint_controller,
            'controllers_file': controllers_file,
            'cartesian_controllers': cartesian_controllers,
            'namespace': namespace,
        }.items(),
        condition=IfCondition(launch_controllers))
   
    moveit = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare(moveit_config_package),
                'launch',
                'moveit.launch.py'
            ])
        ]),
        launch_arguments={
            'robot_description': robot_description_content,
            'moveit_config_package': moveit_config_package,
            'moveit_config_file': moveit_config_file,
            'use_sim_time': use_sim_time,
            'tf_prefix': robot_name,
            'with_gripper': with_gripper_condition,
            'robot_description_semantic': robot_description_semantic_content,
            'namespace': namespace,
        }.items(),
        condition=IfCondition(launch_moveit))

    skills = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('rbs_skill_servers'),
                'launch',
                'skills.launch.py'
            ])
        ]),
        launch_arguments={
            'robot_description': robot_description_content,
            'robot_description_semantic': robot_description_semantic_content,
            'robot_description_kinematics': robot_description_kinematics,
            'use_sim_time': use_sim_time,
            'with_gripper_condition': with_gripper_condition,
            'namespace': namespace
        }.items()
    )
    
    task_planner = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('rbs_task_planner'),
                'launch',
                'task_planner.launch.py'
            ])
        ]),
        launch_arguments={
            # TBD
        }.items(),
        condition=IfCondition(launch_task_planner))
    
    perception = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('rbs_perception'),
                'launch',
                'perception.launch.py'
            ])
        ]),
        launch_arguments={
            # TBD
        }.items(),
        condition=IfCondition(launch_perception))

    nodes_to_start = [
        robot_state_publisher,
        control,
        moveit,
        skills,
        task_planner,
        perception,
        # rviz
    ]
    return nodes_to_start


def generate_launch_description():
    declared_arguments = []
    declared_arguments.append(
        DeclareLaunchArgument(
            "robot_type",
            description="Type of robot by name",
            choices=["rbs_arm","ur3", "ur3e", "ur5", "ur5e", "ur10", "ur10e", "ur16e"],
            default_value="rbs_arm",
        )
    )
    # General arguments
    declared_arguments.append(
        DeclareLaunchArgument(
            "controllers_file",
            default_value="rbs_arm_controllers_gazebosim.yaml",
            description="YAML file with the controllers configuration.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "description_package",
            default_value="rbs_arm",
            description="Description package with robot URDF/XACRO files. Usually the argument \
        is not set, it enables use of a custom description.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "description_file",
            default_value="rbs_arm_modular.xacro",
            description="URDF/XACRO description file with the robot.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "robot_name",
            default_value="arm0",
            description="tf_prefix of the joint names, useful for \
        multi-robot setup. If changed than also joint names in the controllers' configuration \
        have to be updated.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "start_joint_controller",
            default_value="true",
            description="Enable headless mode for robot control",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "initial_joint_controller",
            default_value="joint_trajectory_controller",
            description="Robot controller to start.",
        )
    )
    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(
            "gripper_name",
            default_value="",
            choices=["rbs_gripper", ""],
            description="choose gripper by name (leave empty if hasn't)",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument("with_gripper", 
                              default_value="true", 
                              description="With gripper or not?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("launch_rviz", 
                              default_value="true", 
                              description="Launch RViz?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("sim_gazebo", 
                              default_value="true", 
                              description="Gazebo Simulation")
    )
    declared_arguments.append(
        DeclareLaunchArgument("env_manager", 
                              default_value="true", 
                              description="Launch env_manager?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("launch_sim", 
                              default_value="true", 
                              description="Launch simulator (Gazebo)?\
                              Most general arg")
    )
    declared_arguments.append(
        DeclareLaunchArgument("launch_moveit", 
                              default_value="true", 
                              description="Launch moveit?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("launch_perception", 
                              default_value="false", 
                              description="Launch perception?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("launch_task_planner", 
                              default_value="false", 
                              description="Launch task_planner?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("cartesian_controllers", 
                              default_value="false", 
                              description="Load cartesian\
                              controllers?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("hardware", 
                              choices=["gazebo", "mock"],
                              default_value="gazebo", 
                              description="Choose your harware_interface")
    )
    declared_arguments.append(
        DeclareLaunchArgument("launch_controllers", 
                              default_value="true", 
                              description="Launch controllers?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("gazebo_gui", 
                              default_value="false", 
                              description="Launch gazebo with gui?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("namespace", 
                              default_value="", 
                              description="The ROS2 namespace of a robot")
    )
    declared_arguments.append(
        DeclareLaunchArgument("x", 
                              default_value="0.0", 
                              description="Position of robot in world by X")
    )
    declared_arguments.append(
        DeclareLaunchArgument("y", 
                              default_value="0.0", 
                              description="Position of robot in world by Y")
    )
    declared_arguments.append(
        DeclareLaunchArgument("z", 
                              default_value="0.0", 
                              description="Position of robot in world by Z")
    )
    declared_arguments.append(
        DeclareLaunchArgument("roll", 
                              default_value="0.0", 
                              description="Position of robot in world by Z")
    )
    declared_arguments.append(
        DeclareLaunchArgument("pitch", 
                              default_value="0.0", 
                              description="Position of robot in world by Z")
    )
    declared_arguments.append(
        DeclareLaunchArgument("yaw", 
                              default_value="0.0", 
                              description="Position of robot in world by Z")
    )


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