import os
from launch import LaunchDescription
from launch.actions import (
    IncludeLaunchDescription,
    DeclareLaunchArgument,
    RegisterEventHandler,
    SetEnvironmentVariable,
    TimerAction,
)
from launch.conditions import IfCondition
from launch.event_handlers import OnProcessExit
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import (
    Command,
    EnvironmentVariable,
    LaunchConfiguration,
    PythonExpression,
)
from launch_ros.actions import Node
from launch_ros.parameter_descriptions import ParameterValue
from ament_index_python.packages import get_package_share_directory


def generate_launch_description():
    pkg_gazebo = get_package_share_directory('agv_pro_gazebo')
    pkg_description = get_package_share_directory('agv_pro_description')
    pkg_ros_gz_sim = get_package_share_directory('ros_gz_sim')

    # Shared robot description (agv_pro_description) built for simulation.
    xacro_file = os.path.join(pkg_description, 'urdf', 'agv_pro.urdf.xacro')
    controllers_file = os.path.join(pkg_gazebo, 'config', 'agv_control.yaml')
    world_file = os.path.join(pkg_gazebo, 'worlds', 'empty.world')
    rviz_config = os.path.join(pkg_gazebo, 'rviz', 'agvpro_display.rviz')

    use_sim_time = LaunchConfiguration('use_sim_time')
    use_rviz = LaunchConfiguration('use_rviz')
    headless = LaunchConfiguration('headless')

    robot_description = {
        'robot_description': ParameterValue(
            Command([
                'xacro ', xacro_file,
                ' sim:=true',
                ' controllers_file:=', controllers_file,
            ]),
            value_type=str,
        ),
        'use_sim_time': use_sim_time,
    }

    # Let gz-sim resolve the package:// mesh URIs in the shared description.
    # URDF->SDF conversion rewrites package://agv_pro_description/... to
    # model://agv_pro_description/..., which gz finds by searching
    # GZ_SIM_RESOURCE_PATH for a directory named 'agv_pro_description'.
    gz_resource_path = SetEnvironmentVariable(
        name='GZ_SIM_RESOURCE_PATH',
        value=[
            os.path.dirname(pkg_description),
            ':',
            EnvironmentVariable('GZ_SIM_RESOURCE_PATH', default_value=''),
        ],
    )

    # Start gz-sim (Gazebo Harmonic) with the world.
    # '-s' (server only) is added when headless:=true.
    gz_args = PythonExpression([
        "'-r -v4 -s ' + ", repr(world_file),
        " if '", headless, "' == 'true' ",
        "else '-r -v4 ' + ", repr(world_file),
    ])
    gz_sim = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            os.path.join(pkg_ros_gz_sim, 'launch', 'gz_sim.launch.py')
        ),
        launch_arguments={'gz_args': gz_args}.items(),
    )

    # Publish the robot description (also consumed by gz_ros2_control)
    robot_state_publisher = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        name='robot_state_publisher',
        output='screen',
        parameters=[robot_description],
    )

    # Spawn the robot into gz-sim from the robot_description topic.
    # Delayed so robot_state_publisher has latched /robot_description before the
    # gz_ros2_control plugin (loaded on spawn) reads it — avoids a startup race
    # where the in-sim controller_manager sees no <ros2_control> tag.
    spawn_entity = Node(
        package='ros_gz_sim',
        executable='create',
        arguments=['-topic', 'robot_description', '-name', 'agv_pro', '-z', '0.1'],
        output='screen',
    )
    delayed_spawn = TimerAction(period=3.0, actions=[spawn_entity])

    # Bridge the simulation clock so ROS nodes get /clock
    clock_bridge = Node(
        package='ros_gz_bridge',
        executable='parameter_bridge',
        arguments=['/clock@rosgraph_msgs/msg/Clock[gz.msgs.Clock'],
        output='screen',
    )

    # Controllers (controller_manager runs inside the gz_ros2_control plugin)
    joint_state_broadcaster_spawner = Node(
        package='controller_manager',
        executable='spawner',
        arguments=[
            'joint_state_broadcaster',
            '--controller-manager', '/controller_manager',
        ],
        output='screen',
    )

    # Remap the controller's reference topic (~/reference) to the standard
    # /cmd_vel so teleop and nav2 (which publish geometry_msgs/TwistStamped)
    # drive the robot directly. NOTE: the remap key must be the private name
    # '~/reference' — a bare 'reference:=/cmd_vel' is silently ignored.
    mecanum_drive_controller_spawner = Node(
        package='controller_manager',
        executable='spawner',
        arguments=[
            'mecanum_drive_controller',
            '--controller-manager', '/controller_manager',
            '--controller-ros-args', '-r ~/reference:=/cmd_vel',
        ],
        output='screen',
    )

    rviz = Node(
        package='rviz2',
        executable='rviz2',
        name='rviz2',
        arguments=['-d', rviz_config],
        parameters=[{'use_sim_time': use_sim_time}],
        condition=IfCondition(use_rviz),
        output='screen',
    )

    return LaunchDescription([
        DeclareLaunchArgument(
            'use_sim_time',
            default_value='true',
            description='Use the Gazebo simulation clock',
        ),
        DeclareLaunchArgument(
            'use_rviz',
            default_value='true',
            description='Launch RViz2',
        ),
        DeclareLaunchArgument(
            'headless',
            default_value='false',
            description='Run gz-sim without the GUI (server only)',
        ),
        gz_resource_path,
        gz_sim,
        robot_state_publisher,
        clock_bridge,
        delayed_spawn,
        # Load controllers only after the robot has been spawned
        RegisterEventHandler(
            OnProcessExit(
                target_action=spawn_entity,
                on_exit=[joint_state_broadcaster_spawner],
            )
        ),
        RegisterEventHandler(
            OnProcessExit(
                target_action=joint_state_broadcaster_spawner,
                on_exit=[mecanum_drive_controller_spawner],
            )
        ),
        rviz,
    ])
