import os
from launch import LaunchDescription
from launch.actions import (
    DeclareLaunchArgument,
    IncludeLaunchDescription,
)
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

def generate_launch_description():
    declared_arguments = []
    declared_arguments.append(
        DeclareLaunchArgument(
            "rbs_robot_type",
            description="Type of robot by name",
            choices=["ur3", "ur3e", "ur5", "ur5e", "ur10", "ur10e", "ur16e"],
            default_value="ur5e",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "safety_limits",
            default_value="true",
            description="Enables the safety limits controller if true.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "safety_pos_margin",
            default_value="0.15",
            description="The margin to lower and upper limits in the safety controller.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "safety_k_position",
            default_value="20",
            description="k-position factor in the safety controller.",
        )
    )
    # General arguments
    declared_arguments.append(
        DeclareLaunchArgument(
            "controllers_file",
            default_value="ur_plus_gripper_controllers.yaml",
            description="YAML file with the controllers configuration.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "controllers_with_gripper_file",
            default_value="ur_plus_gripper_controllers.yaml",
            description="YAML file with the UR + gripper_controller configuration.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "description_package",
            default_value="ur_description",
            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="ur.urdf.xacro",
            description="URDF/XACRO description file with the robot.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "prefix",
            default_value='""',
            description="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(
            "initial_gripper_controller",
            default_value="gripper_controller",
            description="Robot controller to start.",
        )
    )

    declared_arguments.append(
        DeclareLaunchArgument(
            "moveit_config_package",
            default_value="ur_moveit_config",
            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="ur.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("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("sim_mujoco", default_value="false", description="Mujoco Simulation")
    )
    declared_arguments.append(
        DeclareLaunchArgument("sim_fake", default_value="false", description="Mock contollers")
    )
    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="true", description="Launch task_planner?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("cartesian_controllers", default_value="false", description="Load cartesian controllers?")
    )

    # Initialize Arguments
    rbs_robot_type = LaunchConfiguration("rbs_robot_type")
    safety_limits = LaunchConfiguration("safety_limits")
    safety_pos_margin = LaunchConfiguration("safety_pos_margin")
    safety_k_position = LaunchConfiguration("safety_k_position")
    # 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")
    prefix = LaunchConfiguration("prefix")
    start_joint_controller = LaunchConfiguration("start_joint_controller")
    initial_joint_controller = LaunchConfiguration("initial_joint_controller")
    launch_rviz = LaunchConfiguration("launch_rviz")
    launch_simulation = LaunchConfiguration("launch_sim")
    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")
    sim_gazebo = LaunchConfiguration("sim_gazebo")
    sim_mujoco = LaunchConfiguration("sim_mujoco")
    sim_fake = LaunchConfiguration("sim_fake")
    env_manager = LaunchConfiguration("env_manager")


    initial_joint_controllers_file_path = PathJoinSubstitution(
        [FindPackageShare(description_package), "config", controllers_file]
    )

    rviz_config_file = PathJoinSubstitution(
        [FindPackageShare("rbs_bringup"), "config", "fork_view.rviz"]
    )
    
    mujoco_model = PathJoinSubstitution(
        [FindPackageShare("rbs_simulation"), "mujoco_model", "current_mj.xml"]
    )
    
    robot_description_content = Command(
        [
            PathJoinSubstitution([FindExecutable(name="xacro")]),
            " ",
            PathJoinSubstitution(
                [FindPackageShare(description_package), "urdf", description_file]
            ),
            " ",
            "safety_limits:=", safety_limits, " ",
            "safety_pos_margin:=", safety_pos_margin, " ",
            "safety_k_position:=", safety_k_position, " ",
            "name:=", "ur", " ",
            "ur_type:=", rbs_robot_type, " ",
            "prefix:=", prefix, " ",
            "sim_mujoco:=", sim_mujoco, " ",
            "sim_gazebo:=", sim_gazebo, " ",
            "sim_fake:=", sim_fake, " ",
            "simulation_controllers:=", initial_joint_controllers_file_path, " ",
            "with_gripper:=", with_gripper_condition, " ",
            "mujoco_model:=", mujoco_model,
        ]
    )
    robot_description = {"robot_description": robot_description_content}
    
    # MoveIt Configuration
    robot_description_semantic_content = Command(
        [
            PathJoinSubstitution([FindExecutable(name="xacro")]),
            " ",
            PathJoinSubstitution(
                [FindPackageShare(moveit_config_package), "srdf", moveit_config_file]
            ), " ",
            "name:=", "ur", " ",
            "prefix:=", prefix, " ",
            "with_gripper:=", with_gripper_condition
        ]
    )
    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",
        output="both",
        parameters=[{"use_sim_time": True}, robot_description],
    )
    
    rviz = Node(
        package="rviz2",
        executable="rviz2",
        name="rviz2",
        output="log",
        arguments=["-d", rviz_config_file],
        parameters=[
            robot_description,
            robot_description_semantic,
            robot_description_kinematics
        ],
        condition=IfCondition(launch_rviz))
    
    control = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('ur_description'),
                '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
        }.items())
    
    simulation = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('rbs_simulation'),
                'launch',
                'simulation.launch.py'
            ])
        ]),
        launch_arguments={
            'sim_gazebo': sim_gazebo,
            'rbs_robot_type': rbs_robot_type,
            'env_manager': env_manager
        }.items(),
        condition=IfCondition(launch_simulation))
    
    moveit = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('ur_moveit_config'),
                '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,
            'prefix': prefix,
            'with_gripper': with_gripper_condition,
            'robot_description_semantic': robot_description_semantic_content
        }.items(),
        condition=IfCondition(launch_moveit))
    
    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,
        rviz,
        control,
        simulation,
        moveit,
        task_planner,
        perception
    ]
    return LaunchDescription(declared_arguments + nodes_to_start)