from ntpath import isfile
import os

import xacro
from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import (
    DeclareLaunchArgument,
    IncludeLaunchDescription,
    OpaqueFunction,
)
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch_ros.actions import Node
from launch_ros.substitutions import FindPackageShare
from rbs_utils.launch import load_xacro_args
from robot_builder.external.ros2_control import ControllerManager
from robot_builder.parser.urdf import URDF_parser
from ros_gz_bridge.actions import RosGzBridge


def launch_setup(context, *args, **kwargs):
    robot_type = LaunchConfiguration("robot_type")
    # General arguments
    with_gripper_condition = LaunchConfiguration("with_gripper")
    description_package = LaunchConfiguration("description_package")
    description_file = LaunchConfiguration("description_file")
    use_moveit = LaunchConfiguration("use_moveit")
    moveit_config_package = LaunchConfiguration("moveit_config_package")
    moveit_config_file = LaunchConfiguration("moveit_config_file")
    use_sim_time = LaunchConfiguration("use_sim_time")
    ee_link_name = LaunchConfiguration("ee_link_name").perform(context)
    base_link_name = LaunchConfiguration("base_link_name").perform(context)
    control_space = LaunchConfiguration("control_space").perform(context)
    control_strategy = LaunchConfiguration("control_strategy").perform(context)
    interactive = LaunchConfiguration("interactive").perform(context)
    real_robot = LaunchConfiguration("real_robot").perform(context)

    use_rbs_utils = LaunchConfiguration("use_rbs_utils")
    assembly_config_name = LaunchConfiguration("assembly_config_name")

    description_package_abs_path = get_package_share_directory(
        description_package.perform(context)
    )

    controllers_file = os.path.join(
        description_package_abs_path, "config", "controllers.yaml"
    )

    xacro_file = os.path.join(
        description_package_abs_path,
        "urdf",
        description_file.perform(context),
    )

    # xacro_config_file = f"{description_package_abs_path}/config/xacro_args.yaml"
    xacro_config_file = os.path.join(
        description_package_abs_path, "urdf", "xacro_args.yaml"
    )

    mappings_data = load_xacro_args(xacro_config_file, locals())

    robot_description_doc = xacro.process_file(xacro_file, mappings=mappings_data)

    robot_description_semantic_content = ""

    if use_moveit.perform(context) == "true":
        package_dir = get_package_share_directory(moveit_config_package.perform(context))
        
        srdf_file = os.path.join(package_dir, "srdf", moveit_config_file.perform(context))

        if srdf_file.endswith(".xacro"):
            srdf_config_file = os.path.join(package_dir, "srdf", "xacro_args.yaml")
            srdf_mappings = load_xacro_args(srdf_config_file, locals())
            robot_description_semantic_content = xacro.process_file(srdf_file, mappings=srdf_mappings)
            robot_description_semantic_content = robot_description_semantic_content.toprettyxml(indent=" ")
        
        elif srdf_file.endswith(".srdf"):
            with open(srdf_file, "r") as file:
                robot_description_semantic_content = file.read()
                
        control_space = "joint"
        control_strategy = "position"
        interactive = "false"

    robot_description_content = robot_description_doc.toprettyxml(indent="  ")
    robot_description = {"robot_description": robot_description_content}

    # Parse robot and configure controller's file for ControllerManager
    robot = URDF_parser.load_string(
        robot_description_content, ee_link_name=ee_link_name
    )
    ControllerManager.save_to_yaml(
        robot, description_package_abs_path, "controllers.yaml"
    )

    rbs_robot_setup = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            [
                PathJoinSubstitution(
                    [FindPackageShare("rbs_bringup"), "launch", "rbs_robot.launch.py"]
                )
            ]
        ),
        launch_arguments={
            "with_gripper": with_gripper_condition,
            "controllers_file": controllers_file,
            "robot_type": robot_type,
            "description_package": description_package,
            "description_file": description_file,
            "robot_name": robot_type,
            "use_moveit": use_moveit,
            "moveit_config_package": moveit_config_package,
            "moveit_config_file": moveit_config_file,
            "use_sim_time": use_sim_time,
            "use_controllers": "true",
            "robot_description": robot_description_content,
            "robot_description_semantic": robot_description_semantic_content,
            "base_link_name": base_link_name,
            "ee_link_name": ee_link_name,
            "control_space": control_space,
            "control_strategy": control_strategy,
            "interactive_control": interactive,
            "use_rbs_utils": use_rbs_utils,
            "assembly_config_name": assembly_config_name,
        }.items(),
    )

    rviz_config_file = os.path.join(description_package_abs_path, "config", "config.rviz")

    rviz_node = Node(
        package="rviz2",
        executable="rviz2",
        arguments=["-d", rviz_config_file],
        parameters=[{"use_sim_time": True}]
    )

    gazebo_world = os.path.join(description_package_abs_path, "world", "default.sdf")

    gazebo = IncludeLaunchDescription(
            PythonLaunchDescriptionSource(
                [os.path.join(get_package_share_directory('ros_gz_sim'),
                            'launch', 'gz_sim.launch.py')]),
            launch_arguments=[('gz_args', [' -r ', gazebo_world])],
    )

    gazebo_spawn_robot = Node(
        package="ros_gz_sim",
        executable="create",
        arguments=[
            "-name", "arm0",
            "-x", "0.0",
            "-z", "0.0",
            "-y", "0.0",
            "-topic", "/robot_description",
        ],
        output="screen",
        parameters=[{"use_sim_time": True}],
    )

    gz_bridge = RosGzBridge(
        bridge_name="ros_gz_bridge",
        config_file=os.path.join(description_package_abs_path, "config", "gz_bridge.yaml")
    )

    grasping_service = Node(
        package="rbs_mill_assist",
        executable="grasping_service.py"
    )

    get_named_pose_service = Node(
        package="rbs_mill_assist",
        executable="get_key_pose_frame.py"
    )

    get_workspace = Node(
        package="rbs_mill_assist",
        executable="get_workspace",
        parameters=[
            {
                "urdf_path": os.path.join(get_package_share_directory("rbs_mill_assist"), "urdf", "current.urdf"),
                "ee_link": ee_link_name,
                "use_sim_time": True,
                "robot_position": [0.0, 0.0, 0.0]
            }
        ]
    )

    publish_ee_pose = Node(
        package="rbs_mill_assist",
        executable="publish_ee_pose_node",
        parameters=[
            {"use_sim_time": use_sim_time},
            {"base_link": base_link_name},
            {"ee_link": ee_link_name}
        ]
    )

    nodes_to_start = [
        rbs_robot_setup,
        # rviz_node,
        gazebo,
        gazebo_spawn_robot,
        gz_bridge,
        grasping_service,
        get_named_pose_service,
        # publish_ee_pose
        # get_workspace
    ]
    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",
                "ar4",
                "ur3",
                "ur3e",
                "ur5",
                "ur5e",
                "ur10",
                "ur10e",
                "ur16e",
            ],
            default_value="rbs_arm",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "description_package",
            default_value="rbs_mill_assist",
            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.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(
            "moveit_config_package",
            default_value="rbs_mill_assist",
            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",
            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="false", description="With gripper or not?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "use_moveit", default_value="true", description="Launch moveit?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "launch_perception", default_value="false", description="Launch perception?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "use_controllers",
            default_value="true",
            description="Launch controllers?",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "scene_config_file",
            default_value="",
            description="Path to a scene configuration file",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "ee_link_name",
            default_value="grasp_point",
            description="End effector name of robot arm",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "base_link_name",
            default_value="base_link",
            description="Base link name if robot arm",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "control_space",
            default_value="task",
            choices=["task", "joint"],
            description="Specify the control space for the robot (e.g., task space).",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "control_strategy",
            default_value="position",
            choices=["position", "velocity", "effort"],
            description="Specify the control strategy (e.g., position control).",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "interactive",
            default_value="true",
            description="Wheter to run the motion_control_handle controller",
        ),
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "real_robot",
            default_value="false",
            description="Wheter to run on the real robot",
        ),
    )

    declared_arguments.append(
        DeclareLaunchArgument(
            "use_rbs_utils",
            default_value="true",
            description="Wheter to use rbs_utils",
        ),
    )

    declared_arguments.append(
        DeclareLaunchArgument(
            "assembly_config_name",
            default_value="",
            description="Assembly config name from rbs_assets_library",
        ),
    )

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