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,
    TimerAction,
)
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 robot_builder.parser.urdf import URDF_parser
from robot_builder.external.ros2_control import ControllerManager
import yaml


def launch_setup(context, *args, **kwargs):
    # Initialize Arguments
    robot_type = LaunchConfiguration("robot_type")
    # General arguments
    with_gripper_condition = LaunchConfiguration("with_gripper")
    description_package = LaunchConfiguration("description_package")
    description_file = LaunchConfiguration("description_file")
    robot_name = LaunchConfiguration("robot_name")
    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")
    hardware = LaunchConfiguration("hardware")
    gripper_name = LaunchConfiguration("gripper_name")
    scene_config_file = LaunchConfiguration("scene_config_file").perform(context)

    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)

    if not scene_config_file == "":
        config_file = {"config_file": scene_config_file}
    else:
        config_file = {}

    description_package_abs_path = get_package_share_directory(
        description_package.perform(context)
    )

    simulation_controllers = 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"

    # TODO: hide this to another place
    # Load xacro_args
    def param_constructor(loader, node, local_vars):
        value = loader.construct_scalar(node)
        return LaunchConfiguration(value).perform(
            local_vars.get("context", "Launch context if not defined")
        )

    def variable_constructor(loader, node, local_vars):
        value = loader.construct_scalar(node)
        return local_vars.get(value, f"Variable '{value}' not found")

    def load_xacro_args(yaml_file, local_vars):
        # Get valut from ros2 argument
        yaml.add_constructor(
            "!param", lambda loader, node: param_constructor(loader, node, local_vars)
        )

        # Get value from local variable in this code
        # The local variable should be initialized before the loader was called
        yaml.add_constructor(
            "!variable",
            lambda loader, node: variable_constructor(loader, node, local_vars),
        )

        with open(yaml_file, "r") as file:
            return yaml.load(file, Loader=yaml.FullLoader)

    mappings_data = load_xacro_args(xacro_config_file, locals())

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

    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="gripper_grasp_point"
    )
    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,
            "gripper_name": gripper_name,
            "controllers_file": simulation_controllers,
            "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,
            "hardware": hardware,
            "use_controllers": "true",
            "robot_description": robot_description_content,
            "base_link_name": base_link_name,
            "ee_link_name": ee_link_name,
            "control_space": control_space,
            "control_strategy": control_strategy,
        }.items(),
    )

    rbs_runtime = Node(
        package="rbs_runtime",
        executable="runtime",
        parameters=[robot_description, config_file, {"use_sim_time": True}],
    )

    clock_bridge = Node(
        package="ros_gz_bridge",
        executable="parameter_bridge",
        arguments=["/clock@rosgraph_msgs/msg/Clock[ignition.msgs.Clock"],
        output="screen",
    )

    delay_robot_control_stack = TimerAction(period=10.0, actions=[rbs_robot_setup])

    nodes_to_start = [rbs_runtime, clock_bridge, delay_robot_control_stack]
    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",
        )
    )
    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(
            "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(
            "use_moveit", default_value="false", description="Launch moveit?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "launch_perception", default_value="false", description="Launch perception?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "hardware",
            choices=["gazebo", "mock"],
            default_value="gazebo",
            description="Choose your harware_interface",
        )
    )
    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="",
            description="End effector name of robot arm",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "base_link_name",
            default_value="",
            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).",
        )
    )

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