#cloud-config

apt:
  sources:
    ros2:
      source: "deb http://packages.ros.org/ros2/ubuntu $RELEASE main"
      keyid: C1CF 6E31 E6BA DE88 68B1 72B4 F42E D6FB AB17 C654

package_update: true
package_upgrade: true
package_reboot_if_required: false # Does not seem to work well with Multipass

packages:
- build-essential
- cmake
- curl
- git
- lsb-release
- openssh-client
- openssh-server
- python-is-python3
- python3
- python3-pip
- python3-venv
- pv
- tig
- unzip
- vim
- wget
- zip
# ROS 2 specific packages
- python3-colcon-common-extensions
- python3-setuptools
- python3-vcstool
- ros-jazzy-ros-base
# Course dependencies (so we do not need to use Rosdep).
# On native install we would like to have GUI utilities such as RViz and RQT,
# but not here. That is why we do not want to use Rosdep.
- ros-jazzy-ros2cli-common-extensions
- ros-jazzy-foxglove-bridge
- ros-jazzy-image-transport
- ros-jazzy-compressed-image-transport
- ros-jazzy-image-transport-plugins
- ros-jazzy-teleop-twist-keyboard 
# assignment_1
- ros-jazzy-rclpy
- ros-jazzy-rosbag2
- ros-jazzy-ros2launch
- ros-jazzy-tf2-ros
# assignment_2
- python3-opencv
- ros-jazzy-cv-bridge
- ros-jazzy-message-filters
- ros-jazzy-rosbag2
- ros-jazzy-tf-transformations
# assignment_3
- ros-jazzy-nav-msgs
- ros-jazzy-rclpy
# assignment_4
# wasp_autonomous_systems_interfaces
- ros-jazzy-builtin-interfaces
- ros-jazzy-controller-manager
- ros-jazzy-control-msgs
- ros-jazzy-diff-drive-controller
- ros-jazzy-joint-state-broadcaster
- ros-jazzy-rclpy
- ros-jazzy-robot-state-publisher
- ros-jazzy-ros2launch
- ros-jazzy-turtlebot3-description
# wasp_autonomous_systems_interfaces
- ros-jazzy-ament-cmake
- ros-jazzy-std-msgs
- ros-jazzy-rosidl-default-generators
- ros-jazzy-rosidl-default-runtime
- ros-jazzy-ament-lint-auto
- ros-jazzy-ament-lint-common
# webots_ros2_driver
- ros-jazzy-ament-cmake
- ros-jazzy-python-cmake-module 
- ros-jazzy-ament-cmake-python
- ros-jazzy-pluginlib
- ros-jazzy-rclcpp
- ros-jazzy-rclpy
- ros-jazzy-sensor-msgs
- ros-jazzy-std-msgs
- ros-jazzy-geometry-msgs
- ros-jazzy-tf2-geometry-msgs
- ros-jazzy-tf2-ros
- ros-jazzy-vision-msgs
- ros-jazzy-tinyxml2-vendor
- libyaml-cpp-dev
- ros-jazzy-ament-lint-auto
- ros-jazzy-ament-lint-common
# webots_ros2_msgs
- ros-jazzy-ament-cmake
- ros-jazzy-rosidl-default-generators
- ros-jazzy-rosidl-default-runtime
- ros-jazzy-builtin-interfaces
- ros-jazzy-geometry-msgs
- ros-jazzy-std-msgs
- ros-jazzy-vision-msgs
- ros-jazzy-ament-lint-auto
- ros-jazzy-ament-lint-common
# webots_ros2_importer
- ros-jazzy-builtin-interfaces
- python3-collada
- python3-lark
- ros-jazzy-xacro
- ros-jazzy-ament-copyright
- python3-pycodestyle
- python3-pil
- python3-numpy
- python3-pytest
# webots_ros2_control
- ros-jazzy-hardware-interface
- ros-jazzy-controller-manager
- ros-jazzy-pluginlib
- ros-jazzy-rclcpp
- ros-jazzy-rclcpp-lifecycle
- ros-jazzy-ros-environment
- ros-jazzy-ament-cmake
- ros-jazzy-ament-lint-auto
- ros-jazzy-ament-lint-common

write_files:
- content: |

    [[ -z "${COLCON_HOME}" ]] && export COLCON_HOME="$HOME/.colcon"

    [[ -z "${XAUTHORITY}" ]] && export XAUTHORITY=$HOME/.Xauthority
    
    # Source Python virtual environment
    if [ -f /home/ubuntu/ros2_ws/venv/bin/activate ]; then
      source /home/ubuntu/ros2_ws/venv/bin/activate
    fi

    # Source ROS
    if [ -f /opt/ros/jazzy/setup.bash ]; then
      source /opt/ros/jazzy/setup.bash
    fi

    if [ -f /usr/share/colcon_argcomplete/hook/colcon-argcomplete.bash ]; then
      source /usr/share/colcon_argcomplete/hook/colcon-argcomplete.bash
    fi

    # Source local ROS workspace
    if [ -f /home/ubuntu/ros2_ws/install/local_setup.bash ]; then
      source /home/ubuntu/ros2_ws/install/local_setup.bash
    fi

    # Alias for building ROS workspace with symbolic links
    alias cbs="colcon build --symlink-install"

    # Make sure that you can execute the newly installed programs by adding your local bin folder to the PATH
    export PATH="/home/ubuntu/.local/bin:$PATH"

    # Make sure ROS only communicates over localhost
    export ROS_AUTOMATIC_DISCOVERY_RANGE=LOCALHOST
    
    # Tell Webots what shared folder to use

    header="┏━━━━━━━━━━━━━━━━━  \e[5m\e[4mPLEASE NOTE\e[24m\e[25m ━━━━━━━━━━━━━━━━━┓"
    filler="┃                                                ┃"
    footer="┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛"
    text_1="┃         Cannot find shared directory.          ┃"
    text_2="┃         Please mount shared directory.         ┃"
    text_3="┃         \033[1;34m\e]8;;https://github.com/KTH-RPL/wasp_autonomous_systems/wiki/Setup-Multipass-VM#create-and-mount-shared-directory\aLink to Wiki for instructions\e]8;;\a\033[1;31m          ┃"
    
    if [ -d /home/ubuntu/shared ]; then
      if [ ! -d /home/ubuntu/shared/webots ]; then
        mkdir /home/ubuntu/shared/webots
      fi
      host_dir=$(mount | grep '/home/ubuntu/shared' | sed -rn 's/:(.*)\son\s\/home\/ubuntu\/shared.*/\1/p')
      if [[ "$host_dir" =~ ^[a-zA-Z]:.* ]]; then
        # Windows host
        host_dir=$(echo "$host_dir" | sed -re 's/\\/\//g' -e 's/(.*):/\/mnt\/\1/')
      fi
      host_dir=$(echo $host_dir | sed -re 's/\s/\\ /g')
      if [ -z "${host_dir}" ]; then
        echo -e "\n\033[1;31m$header\n$filler\n$text_1\n$text_2\n$filler\n$text_3\n$filler\n$footer\033[0m\n"
      fi
      export WEBOTS_SHARED_FOLDER="$host_dir/webots":/home/ubuntu/shared/webots
    else
      echo -e "\n\033[1;31m$header\n$filler\n$text_1\n$text_2\n$filler\n$text_3\n$filler\n$footer\033[0m\n"
    fi

    # Tell where YOLO models are located
    export YOLO_MODELS_DIR="/home/ubuntu/yolo_models"
  path: /home/ubuntu/.bashrc
  owner: ubuntu:ubuntu
  append: true
  defer: true

- content: |
    3f2778a0807ddbdc55665e79253dd15d  ./real_robot/metadata.yaml
    a774f7605c6cc7869a6c04c9aa12d4e5  ./real_robot/real_robot_0.mcap
    cffee353731b9480e1c80e59f986dd52  ./real_robot/real_robot_1.mcap
    2cb5b99129f7c9ffe03084ad1e5f0202  ./real_robot/real_robot_2.mcap
    267e93bf931be3ab565724727378156d  ./real_robot/real_robot_3.mcap
    6ae77d900a9534621e6eb52b00c46994  ./kitti/metadata.yaml
    6981151ffcfdfe71a48f114af2baf30b  ./kitti/kitti_0.mcap
    de88547bb1db966321afabe045cc9c70  ./kitti/kitti_1.mcap
    3ac056f201e2101d9b6e126d797b8036  ./kitti/kitti_2.mcap
    deeefb621ec16759ad260e804764ac8e  ./kitti/kitti_3.mcap
    682341541be9c4585ce0036b36359d6c  ./kitti/kitti_4.mcap
  path: /home/ubuntu/rosbags_md5sums.txt
  owner: ubuntu:ubuntu
  defer: true

- content: |
    #!/bin/bash

    header="┏━━━━━━━━━━━━━━━━━  \e[5m\e[4mPLEASE NOTE\e[24m\e[25m ━━━━━━━━━━━━━━━━━┓"
    filler="┃                                                ┃"
    footer="┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛"
    text_1="┃         Cannot find shared directory.          ┃"
    text_2="┃         Please mount shared directory.         ┃"
    text_3="┃         \033[1;34m\e]8;;https://github.com/KTH-RPL/wasp_autonomous_systems/wiki/Setup-Multipass-VM#create-and-mount-shared-directory\aLink to Wiki for instructions\e]8;;\a\033[1;31m          ┃"
    
    if [ -d /home/ubuntu/shared ]; then
      host_dir=$(mount | grep '/home/ubuntu/shared' | sed -rn 's/:(.*)\son\s\/home\/ubuntu\/shared.*/\1/p')
      if [[ "$host_dir" =~ ^[a-zA-Z]:.* ]]; then
        # Windows host
        host_dir=$(echo "$host_dir" | sed -re 's/\\/\//g' -e 's/(.*):/\/mnt\/\1/')
      fi
      host_dir=$(echo $host_dir | sed -re 's/\s/\\ /g')
    
      if [ -z "${host_dir}" ]; then
        echo -e "\n\033[1;31m$header\n$filler\n$text_1\n$text_2\n$filler\n$text_3\n$filler\n$footer\033[0m\n"
      else
        echo "Extracting rosbags"
        cd /home/ubuntu/shared
        pv /home/ubuntu/rosbags.tar.zst | tar -I "unzstd -T0" -x
        cd rosbags
        echo "Verifying checksums"
        if md5sum -c /home/ubuntu/rosbags_md5sums.txt; then
          echo -e "\033[1;32m✅ Successfully extracted the rosbags! ✅\033[0m"
        else
          echo -e "\033[1;31m❌ Failed to extract the rosbags! ❌\033[0m"
        fi
      fi
    else
      echo -e "\n\033[1;31m$header\n$filler\n$text_1\n$text_2\n$filler\n$text_3\n$filler\n$footer\033[0m\n"
    fi
  path: /home/ubuntu/extract_rosbags.bash
  owner: ubuntu:ubuntu
  permissions: '0755'
  defer: true

- content: |
    Created ~/.config directory:          WAIT ⚠️
    Cloning course repository:            WAIT ⚠️
    Created Python virtual environment:   WAIT ⚠️
    Activated Python virtual environment: WAIT ⚠️
    Installed PIP packages:               WAIT ⚠️
    Downloaded and exported YOLO models:  WAIT ⚠️
    Built ROS workspace:                  WAIT ⚠️
    Installed course Assignment 1:        WAIT ⚠️
    Installed course Assignment 2:        WAIT ⚠️
    Installed course Assignment 3:        WAIT ⚠️
    Installed course Assignment 4:        WAIT ⚠️
    Downloaded rosbags:                   WAIT ⚠️
    Verified downloaded rosbags:          WAIT ⚠️
  path: /home/ubuntu/setup.log
  owner: ubuntu:ubuntu
  defer: true

- content: |
    #!/bin/bash

    function cleanup {
      sed -i "s/WAIT ⚠️/FAIL ❌ \[Never executed\]/g" $log_file
    }

    trap cleanup EXIT

    log_file=$1

    log_info()
    {
      local start=$EPOCHREALTIME

      eval "$@"
      ret=$?

      local end=$EPOCHREALTIME
      local duration=$(echo "scale=2;($end-$start)/1" | bc)

      if [ $ret -eq 0 ]; then
        sed -i "0,/WAIT ⚠️/s//PASS ✅ \[$duration s\]/" $log_file
      else
        sed -i "0,/WAIT ⚠️/s//FAIL ❌ \[$duration s\] \[Error code: $ret\]/" $log_file

        # Terminate the script early
        exit $ret
      fi
    }

    # Download rosbags in background    
    curl -L -C - https://canvas.kth.se/files/9471655/download\?download_frd\=1\&verifier\=4wOMCwu3QHS0duQGTKpn6HyhO9jCPGZSL3RM9kmu --output /home/ubuntu/rosbags.tar.zst &
    ROSBAGS_DOWNLOAD_PID=$!
    
    log_info "mkdir /home/ubuntu/.config"
    
    log_info "git clone --recurse-submodules https://github.com/KTH-RPL/wasp_autonomous_systems.git -b ht25 /home/ubuntu/ros2_ws/src/wasp_autonomous_systems"

    log_info "python3 -m venv /home/ubuntu/ros2_ws/venv --system-site-packages && touch /home/ubuntu/ros2_ws/venv/COLCON_IGNORE"
    
    log_info ". /home/ubuntu/ros2_ws/venv/bin/activate"
    
    log_info "python3 -m pip install --upgrade pip && python3 -m pip install -r /home/ubuntu/ros2_ws/src/wasp_autonomous_systems/requirements.txt"
    
    mkdir /home/ubuntu/yolo_models

    log_info "cd /home/ubuntu/yolo_models && yolo segment export model=yolo11n-seg.pt format=openvino half=True device=cpu dynamic=True && yolo segment export model=yolo11n-seg.pt format=openvino int8=True device=cpu dynamic=True data=coco128-seg.yaml"
    
    log_info ". /opt/ros/jazzy/setup.bash && cd /home/ubuntu/ros2_ws && python3 -m colcon build --cmake-args -DCMAKE_BUILD_TYPE=Release --symlink-install --packages-up-to-regex assignment"

    for i in 1 2 3 4; do
      log_info "test -d '/home/ubuntu/ros2_ws/install/assignment_$i'"
    done

    log_info "wait $ROSBAGS_DOWNLOAD_PID"

    log_info "echo db0c7fe93e4b78649c95a5f5b4e7dec4  /home/ubuntu/rosbags.tar.zst | md5sum -c -"

  path: /home/ubuntu/wasp_setup.bash
  owner: ubuntu:ubuntu
  permissions: '0755'
  defer: true

runcmd:
- [ su, ubuntu, -c, "/home/ubuntu/wasp_setup.bash /home/ubuntu/setup.log" ]
- chown root:root /home/ubuntu/setup.log
- rm /home/ubuntu/wasp_setup.bash