import os
from launch import LaunchDescription
from launch.conditions import IfCondition
from launch_ros.actions import Node, PushRosNamespace
from launch.actions import DeclareLaunchArgument, IncludeLaunchDescription, TimerAction
from launch.substitutions import Command, LaunchConfiguration, PythonExpression
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch_ros.parameter_descriptions import ParameterValue
from ament_index_python.packages import get_package_share_directory

def include_lidar(pkg_name, launch_file, enable_lidar, lidar_type, expected_type):

    return IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            os.path.join(
                get_package_share_directory(pkg_name),
                'launch',
                launch_file
            )
        ),
        condition=IfCondition(
            PythonExpression([
                "'", enable_lidar, "' == 'true' and '",
                lidar_type, "' == '", expected_type, "'"
            ])
        )
    )

def generate_launch_description():

    port_name_arg = LaunchConfiguration('port_name')
    namespace = LaunchConfiguration('namespace')
    lidar_type = LaunchConfiguration('lidar_type')
    enable_lidar = LaunchConfiguration('enable_lidar')

    urdf_file = os.path.join(
        get_package_share_directory('agv_pro_description'),
        'urdf',
        'agv_pro.urdf'
    )

    # Pass both namespace and port_name into the xacro processor so that the
    # <ros2_control> block in the URDF picks up the correct serial port.
    robot_description_content = ParameterValue(
        Command([
            'xacro ',
            urdf_file,
            ' namespace:=',
            PythonExpression(['"', namespace, '" + "/" if "', namespace, '" != "" else ""']),
            ' port_name:=',
            port_name_arg,
        ]),
        value_type=str
    )

    declare_port_name_arg = DeclareLaunchArgument(
        'port_name',
        default_value='/dev/agvpro_controller',
        description='Serial port for the AGV Pro base controller'
    )

    declare_namespace_arg = DeclareLaunchArgument(
        'namespace',
        default_value='',
        description='Namespace for nodes'
    )

    declare_enable_lidar_arg = DeclareLaunchArgument(
        'enable_lidar',
        default_value='true',
        description='Whether to launch lidar drivers'
    )

    declare_lidar_type_arg = DeclareLaunchArgument(
        'lidar_type',
        default_value='n10p',
        description='Lidar type: n10p | mid360 | l2'
    )

    ns_action = PushRosNamespace(namespace)

    ros2_controllers_yaml = os.path.join(
        get_package_share_directory('agv_pro_bringup'),
        'config',
        'ros2_controllers.yaml'
    )

    # controller_manager loads the hardware plugin (serial comms, wheel states)
    # and manages the controllers.
    controller_manager = Node(
        package='controller_manager',
        executable='ros2_control_node',
        parameters=[
            {'robot_description': robot_description_content},
            ros2_controllers_yaml,
        ],
        output='screen',
    )

    robot_state_pub = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        name='robot_state_publisher',
        parameters=[{'robot_description': robot_description_content}],
        output='screen'
    )

    # joint_state_broadcaster publishes /joint_states from the hardware interface.
    # Replaces the old static joint_state_publisher.
    joint_state_broadcaster_spawner = Node(
        package='controller_manager',
        executable='spawner',
        arguments=[
            'joint_state_broadcaster',
            '--controller-manager', 'controller_manager',
        ],
        output='screen',
    )

    # mecanum_drive_controller subscribes to ~/reference (TwistStamped).
    # --controller-ros-args remaps it to /cmd_vel so nav2 and teleop work
    # without extra flags (teleop still needs stamped:=true).
    # Delayed slightly so controller_manager is ready before spawning.
    mecanum_drive_controller_spawner = TimerAction(
        period=2.0,
        actions=[
            Node(
                package='controller_manager',
                executable='spawner',
                arguments=[
                    'mecanum_drive_controller',
                    '--controller-manager', 'controller_manager',
                    '--controller-ros-args', '-r reference:=/cmd_vel',
                ],
                output='screen',
            )
        ]
    )

    # Relay /cmd_vel (TwistStamped) → /mecanum_drive_controller/reference.
    # This bridges the standard nav2/teleop topic to the controller's
    # internal subscription, which ros2_control does not remap at load time.
    # lazy:=false ensures the subscription exists before any publisher appears.
    cmd_vel_relay = Node(
        package='topic_tools',
        executable='relay',
        name='cmd_vel_relay',
        parameters=[{
            'input_topic':  '/cmd_vel',
            'output_topic': '/mecanum_drive_controller/reference',
            'lazy': False,
        }],
        output='screen',
    )

    lidar_launchs = [
        include_lidar('lslidar_driver', 'lsn10p_launch.py', enable_lidar, lidar_type, 'n10p'),
        include_lidar('agv_pro_bringup', 'MID360_launch.py', enable_lidar, lidar_type, 'mid360'),
        include_lidar('agv_pro_bringup', 'unitree_l2_launch.py', enable_lidar, lidar_type, 'l2'),
    ]

    return LaunchDescription(
        [
            declare_port_name_arg,
            declare_namespace_arg,
            declare_enable_lidar_arg,
            declare_lidar_type_arg,
            ns_action,
            controller_manager,
            robot_state_pub,
            joint_state_broadcaster_spawner,
            mecanum_drive_controller_spawner,
            cmd_vel_relay,
            *lidar_launchs,
        ]
    )