import os
import xacro
from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import IncludeLaunchDescription, RegisterEventHandler, TimerAction
from launch.event_handlers import OnProcessExit
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch_ros.actions import Node

# Absolute and relative path variables
world_path = './src/world_description/corn_field.sdf'
arm_urdf_path = '/home/marcin/arm_ws/robotarm_description/urdf/robot_arm.urdf'
rover_sdf_path = '/home/marcin/arm_ws/src/world_description/harvester.sdf'

def generate_launch_description():
    moveit_config_pkg = get_package_share_directory('arm_moveit_config')
    
    # 1. Include the standard MoveIt 2 move_group launch file (Delayed for stability)
    move_group_launch = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            os.path.join(moveit_config_pkg, 'launch', 'move_group.launch.py')
        )
    )

    # 2. Launch Gazebo Sim (Modern Gazebo) loaded directly into your field world
    gazebo_launch = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            os.path.join(get_package_share_directory('ros_gz_sim'), 'launch', 'gz_sim.launch.py')
        ),
        launch_arguments={'gz_args': '-r ' + world_path }.items(),
    )

    # 3. MODIFIED: Spawns the ENTIRE rover model from the master SDF file.
    # Because your SDF includes the arm internally, Gazebo builds the combined asset together.
    spawn_complete_rover = Node(
        package='ros_gz_sim',
        executable='create',
        arguments=[
            '-file', rover_sdf_path,
            '-name', 'harvester_robot', # Matches the model namespace inside your world bridges
            '-z', '0.1'                 # Clear the field ground mesh on entry
        ],
        output='screen',
    )

    # 4. Define ros2_control Spawners
    joint_state_broadcaster_spawner = Node(
        package="controller_manager",
        executable="spawner",
        arguments=["joint_state_broadcaster", "--controller-manager", "/controller_manager"],
    )

    arm_controller_spawner = Node(
        package="controller_manager",
        executable="spawner",
        arguments=["arm_controller", "--controller-manager", "/controller_manager"],
    )

    # 5. Bridge Gazebo Clock and Data Channels to ROS 2 Jazzy
    bridge = Node(
        package='ros_gz_bridge',
        executable='parameter_bridge',
        arguments=[
            '/clock@rosgraph_msgs/msg/Clock[gz.msgs.Clock',
            '/world/corn_field_wide/model/harvester_robot/joint_state@sensor_msgs/msg/JointState[gz.msgs.Model',
            '/arm_camera/image_raw@sensor_msgs/msg/Image[gz.msgs.Image',
            # '/model/harvester_robot/joint/joint_1/cmd_pos@std_msgs/msg/Float64]gz.msgs.Double',
            # 'cmd_vel@geometry_msgs/msg/Twist]gz.msgs.Twist',
            # 'model/harvester_robot/odometry@nav_msgs/msg/Odometry[gz.msgs.Odometry',
            # 'lidar_left@sensor_msgs/msg/LaserScan[gz.msgs.LaserScan',
            # 'lidar_right@sensor_msgs/msg/LaserScan[gz.msgs.LaserScan',
            ],
        remappings=[
            ('/model/harvester_robot/odometry', '/odom')
        ],
        output='screen'
    )

    # 6. Process the xacro file dynamically for the ROS 2 Kinematics Stack
    xacro_file = os.path.join(
        get_package_share_directory('robotarm_description'),
        'urdf',
        'robot_arm.urdf.xacro'
    )
    robot_description_raw = xacro.process_file(xacro_file).toxml()

    # 7. Robot State Publisher (Feeds the mathematical link tree strictly to MoveIt 2)
    robot_state_publisher = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        output='both',
        parameters=[{'robot_description': robot_description_raw, 'use_sim_time': True}]
    )

    # 8. Timing/Order Sequence Event Handlers
    # Waits for the COMPLETE rover to exist before spinning up hardware controller instances
    delay_joint_state_broadcaster = RegisterEventHandler(
        event_handler=OnProcessExit(
            target_action=spawn_complete_rover,
            on_exit=[joint_state_broadcaster_spawner],
        )
    )

    delay_arm_controller = RegisterEventHandler(
        event_handler=OnProcessExit(
            target_action=joint_state_broadcaster_spawner,
            on_exit=[arm_controller_spawner],
        )
    )
    
    delay_move_group_launch = TimerAction(
        period=5.0,
        actions=[move_group_launch]
    )

    return LaunchDescription([
        gazebo_launch,
        spawn_complete_rover, # Spawns the single, complete vehicle
        robot_state_publisher,
        bridge,
        delay_joint_state_broadcaster,
        delay_arm_controller,
        delay_move_group_launch,
    ])