from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument
from launch.conditions import IfCondition, UnlessCondition 
from launch.substitutions import Command, LaunchConfiguration, IfElseSubstitution
from launch_ros.actions import Node, LifecycleNode
from launch_ros.substitutions import FindPackageShare
import os

def generate_launch_description():
    pkg_share = FindPackageShare("").find('mg_control')
    params = os.path.join(pkg_share, 'assets', 'toid_general_params.yaml')
    default_rviz_config_path = os.path.join(pkg_share, 'rviz', 'ros_control.rviz')
    

    description_pkg_share = FindPackageShare("").find('toid_bot_description')
    default_model_path = os.path.join(
        description_pkg_share, 
        'src', 
        'toid_bot_description.urdf'
    )

    visualize = LaunchConfiguration("visualize")

    visualize_arg = DeclareLaunchArgument(
        'visualize',
        default_value='False',
        description="Whether to launch rviz2" 
    )

    use_mock = LaunchConfiguration("use_mock")

    use_mock_arg = DeclareLaunchArgument(
        'use_mock',
        default_value='True',
        description="Whether to use mock controller" 
    )

    odom_broadcast = Node(
        package='tf2_ros',
        executable='static_transform_publisher',
        name='map_to_odom_broadcaster',
        arguments=['0', '0', '0', '0', '0', '0', 'map', 'odom'],
        condition=IfCondition(LaunchConfiguration('visualize'))
    )

    robot_state_publisher = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        name='robot_state_publisher',
        output='screen',
        parameters=[{'robot_description': Command(['xacro ', default_model_path, ' use_mock:=', use_mock])}]
    )

    controller_manager = Node(
        package='controller_manager',
        executable='ros2_control_node',
        output='screen',
        parameters=[params]
    )

    joint_state_broadcaster = Node(
        package='controller_manager',
        executable='spawner',
        output='screen',
        arguments=["joint_state_broadcaster"]
    )

    diffbot_base_controller = Node(
        package='controller_manager',
        executable='spawner',
        output='both',
        arguments=[
            "diffdrive_controller",
            "-p",
            params,
            "--controller-ros-args",
            "-r diffdrive_controller/cmd_vel:=/cmd_vel",
            "--controller-ros-args",
            "-r diffdrive_controller/odom:=/odom",
            "--controller-ros-args",
            IfElseSubstitution(use_mock, 
                "--param odom_frame_id:=odom", 
                "--param odom_frame_id:=odom_expected"
            ),
            "--controller-ros-args",
            IfElseSubstitution(use_mock, 
                "--param enable_odom_tf:=true", 
                "--param enable_odom_tf:=false"
            ),
        ]
    )

    rviz_node = Node(
        package='rviz2',
        executable='rviz2',
        name='rviz2',
        output='screen',
        arguments=['-d', default_rviz_config_path],
        condition=IfCondition(visualize)
    )

    return LaunchDescription([
        visualize_arg,
        use_mock_arg,
        odom_broadcast,
        robot_state_publisher,
        controller_manager,
        joint_state_broadcaster,
        diffbot_base_controller,
        rviz_node
    ])