from launch import LaunchDescription
from launch.actions import (
    DeclareLaunchArgument,
    IncludeLaunchDescription,
    OpaqueFunction,
)
from launch.conditions import IfCondition
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch_ros.substitutions import FindPackageShare
from launch_ros.actions import Node
from nav2_common.launch import RewrittenYaml
import os
from ament_index_python import get_package_share_directory
from rbs_launch_utils.merged_yaml import MergedYaml
from rbs_launch_utils.launch_common import load_yaml
import yaml

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_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")
    hardware = LaunchConfiguration("hardware")
    env_manager = LaunchConfiguration("env_manager")
    launch_controllers = LaunchConfiguration("launch_controllers")
    gazebo_gui = LaunchConfiguration("gazebo_gui")
    gripper_name = LaunchConfiguration("gripper_name")
    robots_config_file = LaunchConfiguration("robots_config_file")
    gazebo_world_filename = LaunchConfiguration("gazebo_world_filename")

    ld = []
    ld.append(IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('rbs_simulation'),
                'launch',
                'simulation_gazebo.launch.py'
            ])
        ]),
        launch_arguments={
            'sim_gazebo': sim_gazebo,
            'debugger': "false",
            'launch_env_manager': "false",
            "gazebo_world_filename": gazebo_world_filename
        }.items(),
        condition=IfCondition(launch_simulation)
    ))
    scene_file = robots_config_file.perform(context)
    robots = load_yaml("rbs_bringup", "config/" + scene_file)

    description_package = description_package.perform(context)
    controllers_file = controllers_file.perform(context)
    config = MergedYaml(context, 
                         os.path.join(get_package_share_directory(description_package), "config", controllers_file),
                         root_keys=[i["name"] for i in robots["scene_config"]],
                         param_rewrites={},
                         convert_types=False,
                         ).merge_yamls()
    for robot in robots["scene_config"]:
        namespace = "/" + robot["name"]
        ld.append(IncludeLaunchDescription(
            PythonLaunchDescriptionSource([
                PathJoinSubstitution([
                    FindPackageShare('rbs_bringup'),
                    "launch",
                    "rbs_robot.launch.py"
                ])
            ]),
            launch_arguments={
                "env_manager": env_manager,
                "with_gripper": with_gripper_condition,
                "gripper_name": gripper_name,
                "controllers_file": controllers_file,
                "robot_type": robot["type"],
                "controllers_file": config,
                "cartesian_controllers": cartesian_controllers,
                "description_package": description_package,
                "description_file": description_file,
                "robot_name": robot["name"],
                "start_joint_controller": start_joint_controller,
                "initial_joint_controller": initial_joint_controller,
                "launch_simulation": launch_simulation,
                "launch_moveit": launch_moveit,
                "launch_task_planner": launch_task_planner,
                "launch_perception": launch_perception,
                "moveit_config_package": moveit_config_package,
                "moveit_config_file": moveit_config_file,
                "use_sim_time": use_sim_time,
                "sim_gazebo": sim_gazebo,
                "hardware": hardware,
                "launch_controllers": launch_controllers,
                "gazebo_gui": gazebo_gui,
                "x": str(robot["pose"]["position"]["x"]),
                "y": str(robot["pose"]["position"]["y"]),
                "z": str(robot["pose"]["position"]["z"])
            }.items()
        ))

        gz_spawner = Node(
            package='ros_gz_sim',
            executable='create',
            arguments=[
                '-name', robot_name,
                '-x', "0",#str(robot["spawn_point"][0]),
                '-y', "0",#str(robot["spawn_point"][1]),
                '-z', "0",#str(robot["spawn_point"][2]),
                '-topic', namespace + '/robot_description'],
            output='screen',
            condition=IfCondition(sim_gazebo))
        ld.append(gz_spawner)

    return ld


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="Name for robot, used to apply namespace for specific robot in multirobot setup",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "start_joint_controller",
            default_value="false",
            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="rbs_gripper",
            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("sim_gazebo", 
                              default_value="true", 
                              description="Gazebo Simulation")
    )
    declared_arguments.append(
        DeclareLaunchArgument("env_manager", 
                              default_value="false", 
                              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="false", 
                              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="true", 
                              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="true", 
                              description="Launch gazebo with gui?")
    )
    declared_arguments.append(
        DeclareLaunchArgument("gazebo_world_filename", 
                              default_value="asm2.sdf", 
                              description="Filename of Gazebo world file to launch")
    )
    declared_arguments.append(
        DeclareLaunchArgument("robots_config_file", 
                              default_value="robot_scene", 
                              description="Filename for config file with robots in scene")
    )

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