#!/bin/bash

#  Copyright (C) 2018-2020 LEIDOS.
# 
#  Licensed under the Apache License, Version 2.0 (the "License"); you may not
#  use this file except in compliance with the License. You may obtain a copy of
#  the License at
# 
#  http://www.apache.org/licenses/LICENSE-2.0
# 
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#  License for the specific language governing permissions and limitations under
#  the License.

# Code below largely based on template from Stack Overflow:
# https://stackoverflow.com/questions/37257551/defining-subcommands-that-take-arguments-in-bash
# Question asked by user 
# DiogoSaraiva (https://stackoverflow.com/users/4465820/diogosaraiva)
# and answered by user 
# Charles Duffy (https://stackoverflow.com/users/14122/charles-duffy)
# Attribution here is in line with Stack Overflow's Attribution policy cc-by-sa found here:
# https://stackoverflow.blog/2009/06/25/attribution-required/


__pull_newest_carma_base() {
    if [ "$1" = "-d" ]; then
        local remote_image=$(wget -q https://registry.hub.docker.com/v1/repositories/usdotfhwastoldev/carma-base/tags -O -  | \
        sed -e 's/[][]//g' -e 's/"//g' -e 's/ //g' | tr '}' '\n'  | \
        awk -F: '{print "usdotfhwastoldev/carma-base:"$3}' | tail -n 1)
    elif [ "$1" = "-c" ]; then
        local remote_image=$(wget -q https://registry.hub.docker.com/v1/repositories/usdotfhwastolcandidate/carma-base/tags -O -  | \
        sed -e 's/[][]//g' -e 's/"//g' -e 's/ //g' | tr '}' '\n'  | \
        awk -F: '{print "usdotfhwastolcandidate/carma-base:"$3}' | tail -n 1)
    else
        local remote_image=$(wget -q https://registry.hub.docker.com/v1/repositories/usdotfhwastol/carma-base/tags -O -  | \
        sed -e 's/[][]//g' -e 's/"//g' -e 's/ //g' | tr '}' '\n'  | \
        awk -F: '{print "usdotfhwastol/carma-base:"$3}' | tail -n 1)
    fi
    
    echo "No local carma-base image found, pulling down the most recent from Dockerhub..."
    docker pull $remote_image
}

##
# Helper function to print to stderr.
# This is useful for allowing bash methods to return strings while still communicating information to the user
##
echoerr() { echo "$@" 1>&2; }

##
# Method returns a fully specified docker image name and tag based on the currently set carma-config.
# The method will return the first tag that appears in the config for the possible organizations and images.
# If the config is not set or an image cannot be found the method will return a carma-base image as the default.
# The user can specify optional organizations or images as sed compatable regex lists.
# For example calling 
#  $(__get_image_from_config 'carma-platform:\|carma-messenger:' 'usdotfhwastol\|usdotfhwastoldev\|usdotfhwastolcandidate')
# 
#  If no parameters are provided the defaults will be used
##
__get_image_from_config() {
    local TARGET_IMAGE="usdotfhwastol/carma-base:latest"
    local SUPPORTED_ORANIZATIONS="usdotfhwastol\|usdotfhwastoldev\|usdotfhwastolcandidate"
    local SUPPORTED_IMAGES="carma-platform:\|carma-messenger:"

    if [ -n "$1" ]
    then
        local SUPPORTED_IMAGES="$1"
    fi

    if [ -n "$2" ]
    then
        
        local SUPPORTED_ORANIZATIONS="$2"
    fi

    # Check if carma-config is set. If so, use what was in the carma config
    if docker container inspect carma-config > /dev/null 2>&1; then

        local IMAGE_BACKUP=$TARGET_IMAGE

        # sed command finds the first instance of an image which is part of the set defined in SUPPORTED_ORGANIZATIONS and SUPPORTED_IMAGES.
        # then returns the name of the image.  
        local TARGET_IMAGE=$(docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
            "sed -n -e '/image:\s*\($SUPPORTED_ORANIZATIONS\)\/\($SUPPORTED_IMAGES\)/ s/.*\image: *//p' /opt/carma/vehicle/config/docker-compose.yml | grep -m 1 '.*'")
            

        # If the sed resulted in an empty string then reset the tag
        if [ -z "$TARGET_IMAGE" ]; then
            local TARGET_IMAGE=$IMAGE_BACKUP
            echoerr "Could not identify usable image from config so defaulting to image: $TARGET_IMAGE"
        fi
    else
        echoerr "No config detected or target image specified so using default: $TARGET_IMAGE"
    fi

    echo "$TARGET_IMAGE"
}

__get_most_recent_carma_base() {
    if [ "$1" = "-d" ]; then
        local USERNAME=usdotfhwastoldev
    elif [ "$1" = "-c" ]; then
        local USERNAME=usdotfhwastolcandidate
    else
        local USERNAME=usdotfhwastol
    fi
    docker image ls --format "{{.Repository}}:{{.Tag}}" "$USERNAME/carma-base" | grep -v "<none>$" | head -n 1
}

carma-config__set() {
    if [ "$1" = "-d" ]; then
        local USERNAME=usdotfhwastoldev
    elif [ "$1" = "-c" ]; then
        local USERNAME=usdotfhwastolcandidate
    else
        local USERNAME=usdotfhwastol
    fi
    if [[ -z $1 ]]; then
        echo "Please specify a tag string for carma-config to set."
        echo "Done."
        exit -1
    fi

    local IMAGE_NAME="$USERNAME/carma-config:$1"
    if [ ! -z "$(echo $1 | grep :)" ]; then
        IMAGE_NAME=$1
    fi

    if docker container inspect carma-config > /dev/null 2>&1; then
        echo "Clearing existing CARMA configuration instance..."
        carma__stop
        echo "Deleting old CARMA config..."
        docker rm carma-config
    fi

    echo "Setting $IMAGE_NAME as current CARMA configuration instance..."
    docker run --name carma-config "$IMAGE_NAME"
}

carma__attach() {
    local CARMA_DOCKER_FILE="`docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c 'cat /opt/carma/vehicle/config/docker-compose.yml'`"

    echo "Attaching to CARMA container STDOUT..."
    echo "$CARMA_DOCKER_FILE" | docker-compose -p carma -f - logs --follow --tail=10
}

carma-config__edit() {
    if ! docker container inspect carma-config > /dev/null 2>&1; then
        echo "No existing CARMA configuration found, nothing to edit. Please set a config."
        echo "Done."
        exit 1
    fi

    echo "Opening shell inside carma-config container with read/write privileges..."

    local carma_base=$(__get_image_from_config carma-base:)
    if [[ -z $carma_base ]]; then
        __pull_newest_carma_base
        carma_base=$(__get_most_recent_carma_base)
    fi

    docker run -it --rm --volumes-from carma-config $carma_base bash
}

carma-config__inspect() {
    if ! docker container inspect carma-config > /dev/null 2>&1; then
        echo "No existing CARMA configuration found, nothing to inspect. Please set a config."
        echo "Done."
        exit 1
    fi

    echo "Opening shell inside carma-config container with read-only privileges..."

    local carma_base=$(__get_most_recent_carma_base)
    if [[ -z $carma_base ]]; then
        __pull_newest_carma_base
        carma_base=$(__get_most_recent_carma_base)
    fi

    docker run -it --rm --volumes-from carma-config:ro $carma_base bash
}

carma-config__reset() {
    if ! docker container inspect carma-config > /dev/null 2>&1; then
        echo "No existing CARMA configuration found, nothing to reset. Please set a config."
        echo "Done."
        exit 1
    fi

    local CURRENT_IMAGE=`docker container inspect --format='{{.Config.Image}}' carma-config`
    echo "Found current config: $CURRENT_IMAGE, resetting to base state"
    carma-config__set $CURRENT_IMAGE
}

carma-config__list_local() {
    echo "Locally installed images: "
    echo ""
    echo "usdotfhwastol images:"
    docker images usdotfhwastol/carma-config
    echo ""
    echo "usdotfhwastoldev images:"
    docker images usdotfhwastoldev/carma-config
    echo ""
    echo "usdotfhwastolcandidate images:"
    docker images usdotfwhastolcandidate/carma-config
}

carma-config__list_remote() {
    if [ "$1" = "-d" ]; then
        wget -q https://registry.hub.docker.com/v1/repositories/usdotfhwastoldev/carma-config/tags -O -  | \
        sed -e 's/[][]//g' -e 's/"//g' -e 's/ //g' | tr '}' '\n'  | \
        awk -F: 'BEGIN {print "Remotely available images from usdotfhwastoldev Dockerhub:\nIMAGE\t\t\t\tTAG"} {print "usdotfhwastoldev/carma-config\t"$3}'
    elif [ "$1" = "-c" ]; then
        wget -q https://registry.hub.docker.com/v1/repositories/usdotfhwastolcandidate/carma-config/tags -O -  | \
        sed -e 's/[][]//g' -e 's/"//g' -e 's/ //g' | tr '}' '\n'  | \
        awk -F: 'BEGIN {print "Remotely available images from usdotfhwastolcandidate Dockerhub:\nIMAGE\t\t\t\tTAG"} {print "usdotfhwastolcandidate/carma-config\t"$3}' 
    else
        wget -q https://registry.hub.docker.com/v1/repositories/usdotfhwastol/carma-config/tags -O -  | \
        sed -e 's/[][]//g' -e 's/"//g' -e 's/ //g' | tr '}' '\n'  | \
        awk -F: 'BEGIN {print "Remotely available images from usdotfhwastol Dockerhub:\nIMAGE\t\t\t\tTAG"} {print "usdotfhwastol/carma-config\t"$3}'    
    fi
}

carma-config__install() {
    if [ "$1" = "-d" ]; then
        local USERNAME=usdotfhwastoldev
        shift
    elif [ "$1" = "-c" ]; then
        local USERNAME=usdotfhwastolcandidate
        shift
    else
        local USERNAME=usdotfhwastol
    fi

    local IMAGE_NAME="$USERNAME/carma-config:$1"
    if [ ! -z "$(echo $1 | grep :)" ]; then
        IMAGE_NAME=$1
    fi

    echo "Downloading $IMAGE_NAME..."
    docker pull $IMAGE_NAME
    echo "Building temporary container of $IMAGE_NAME to read dependency data..."

    if docker container inspect carma-config-tmp > /dev/null 2>&1; then
        echo "Cleaning up temporary containers from previous install..."
        docker rm carma-config-tmp
    fi
    docker run --name carma-config-tmp $IMAGE_NAME

    local CARMA_DOCKER_FILE="`docker run --rm --volumes-from carma-config-tmp:ro --entrypoint sh busybox:latest -c 'cat /opt/carma/vehicle/config/docker-compose.yml' | sed s/carma-config/carma-config-tmp/g`"
    local CARMA_BACKGROUND_DOCKER_FILE="`docker run --rm --volumes-from carma-config-tmp:ro --entrypoint sh busybox:latest -c 'cat /opt/carma/vehicle/config/docker-compose-background.yml' | sed s/carma-config/carma-config-tmp/g`"

    echo "Downloading $IMAGE_NAME dependencies..."
    echo "$CARMA_DOCKER_FILE" | docker-compose -f - pull
    echo "$CARMA_BACKGROUND_DOCKER_FILE" | docker-compose -f - pull

    echo "Cleaning up temporary container..."
    docker rm carma-config-tmp
}

carma__start() {
    if ! docker container inspect carma-config > /dev/null 2>&1; then
        echo "No existing CARMA configuration found, nothing to start. Please set a config."
        echo "Done."
        exit 1
    fi

    echo "Starting CARMA background processes..."
    docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
    'cat /opt/carma/vehicle/config/docker-compose-background.yml' | \
    docker-compose -f - -p carma-background up -d

    if [ "$1" = "all" ]; then 
        shift
        echo "Starting CARMA Platform foreground processes..."
        docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
        'cat /opt/carma/vehicle/config/docker-compose.yml' | \
        docker-compose -f - -p carma up $@
    elif [ ! -z "$1" ]; then
        echo "Unrecognized argument \"start $1\""
    fi
}

carma__stop() {
    if ! docker container inspect carma-config > /dev/null 2>&1; then
        echo "No existing CARMA configuration found, nothing to stop. Please set a config."
        echo "Done."
        exit 1
    fi

    echo "Shutting down CARMA processes..."
    docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
    'cat /opt/carma/vehicle/config/docker-compose.yml' | \
    docker-compose -f - -p carma down

    if [ "$1" = "all" ]; then 
        echo "Shutting down CARMA background processes..."
        docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
        'cat /opt/carma/vehicle/config/docker-compose-background.yml' | \
        docker-compose -f - -p carma-background down
    elif [ ! -z "$1" ]; then
        echo "Unrecognized argument \"stop $1\""
    fi
}

carma__ps() {
    if ! docker container inspect carma-config > /dev/null 2>&1; then
        echo "No existing CARMA configuration found, nothing to report. Please set a config."
        echo "Done."
        exit 1
    fi

    echo "CARMA Background Processes:"
    docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
    'cat /opt/carma/vehicle/config/docker-compose-background.yml' | \
    docker-compose -f - -p carma-background ps

    echo "CARMA Foreground Processes:"
    docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
    'cat /opt/carma/vehicle/config/docker-compose.yml' | \
    docker-compose -f - -p carma ps
}

carma-config__status() {
    if ! docker container inspect carma-config > /dev/null 2>&1; then
        echo "No existing CARMA configuration found, nothing to report. Please set a config."
        echo "Done."
        exit 1
    fi

    if [ -z "$1" ]; then
        local CURRENT_IMAGE=`docker container inspect --format='{{.Config.Image}}' carma-config`
        echo "Current configuration is loaded from image: $CURRENT_IMAGE"
        echo ""
        echo "  -- docker-compose.yml:"
        docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
        'cat /opt/carma/vehicle/config/docker-compose.yml'
        echo "  -- docker-compose-background.yml:"
        docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
        'cat /opt/carma/vehicle/config/docker-compose-background.yml'
        echo "  -- carma.launch:"
        docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
        'cat /opt/carma/vehicle/config/carma.launch'
        echo "  -- carma.config.js:"
        docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
        'cat /opt/carma/vehicle/config/carma.config.js'
        echo "  -- carma.urdf:"
        docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
        'cat /opt/carma/vehicle/config/carma.urdf'
    else
        docker run --rm --volumes-from carma-config:ro --entrypoint sh busybox:latest -c \
        "cat /opt/carma/vehicle/config/$1"
    fi
}

carma__exec() {
    local USERNAME=usdotfhwastol
    local TAG=latest
    
    local TARGET_IMAGE="$USERNAME/carma-platform:$TAG"

    # Check if the -i argument was provided
    # If the image was specified then use that

    if [ "$1" = "--gui" ]; then 
        local RUNTIME=nvidia
        echo "setting docker runtime to $RUNTIME"
        shift
    fi

    if [ "$1" = "-i" ]; then 

        if [ -z "$2" ]; then
            echo " When -i is specified you must enter a carma docker image and tag."
            exit -1
        fi

        local TARGET_IMAGE="$2"

        shift
        shift
    # If the image was not specified try to use what was in the carma config
    else
        local TARGET_IMAGE=$(__get_image_from_config 'carma-platform:\|carma-messenger:')
    fi

    # Verify target image exists
    if [ -z "$(docker images -q $TARGET_IMAGE 2> /dev/null)" ]; then
        echo "Specified image could not be found. carma exec attempted to use image $TARGET_IMAGE but it was not found."
        echo "If an alternative image is available, please specify it using the -i option."
        exit -1
    fi

    # Docker setup here is based on Pierre Killy's docker-ros-box https://github.com/pierrekilly/docker-ros-box located here.
    # Docker ROS Box is licensed under the MIT licesne a copy of which can be found in the carma-platform repository here https://github.com/usdot-fhwa-stol/carma-platform/blob/develop/docs/License.md

    local XSOCK=/tmp/.X11-unix
    local XAUTH=/tmp/.docker.xauth
    touch $XAUTH
    xauth nlist $DISPLAY | sed -e 's/^..../ffff/' | xauth -f $XAUTH nmerge -

    echo "Trying to open shell inside image: $TARGET_IMAGE"

    # If a command was specified then assign it otherwise use bash
    local entrypoint="/bin/bash"
    local INIT_FILE="/home/carma/.base-image/init-env.sh"

    if [ -n "$1" ]; then

        local COMMAND="source $INIT_FILE && ${@}"
        
        docker run \
            -e DISPLAY=$DISPLAY \
            --runtime=$RUNTIME \
            --group-add video \
            --volume=$XSOCK:$XSOCK:rw \
            --volume=$XAUTH:$XAUTH:rw \
            --volume=/opt/carma/logs:/opt/carma/logs \
            --volume=/opt/carma/.ros:/opt/carma/.ros \
            --volume=/opt/carma/vehicle/calibration:/opt/carma/vehicle/calibration \
            --volume=/opt/carma/yolo:/opt/carma/yolo \
            --volume=/opt/carma/maps:/opt/carma/maps \
            --volume=/opt/carma/routes:/opt/carma/routes \
            --env="XAUTHORITY=${XAUTH}" \
            --env QT_X11_NO_MITSHM=1 \
            --device=/dev/dri:/dev/dri \
            --network=host \
            --entrypoint="$entrypoint" \
            -it $TARGET_IMAGE -c "$COMMAND"

    else
        docker run \
            -e DISPLAY=$DISPLAY \
            --group-add video \
            --volume=$XSOCK:$XSOCK:rw \
            --volume=$XAUTH:$XAUTH:rw \
            --volume=/opt/carma/logs:/opt/carma/logs \
            --volume=/opt/carma/.ros:/opt/carma/.ros \
            --volume=/opt/carma/vehicle/calibration:/opt/carma/vehicle/calibration \
            --volume=/opt/carma/yolo:/opt/carma/yolo \
            --volume=/opt/carma/maps:/opt/carma/maps \
            --volume=/opt/carma/routes:/opt/carma/routes \
            --env="XAUTHORITY=${XAUTH}" \
            --device=/dev/dri:/dev/dri \
            --network=host \
            --entrypoint="$entrypoint" \
            -it $TARGET_IMAGE --init-file "$INIT_FILE"
    fi
}

carma__help() {
    cat <<HELP
-------------------------------------------------------------------------------
| USDOT FHWA STOL CARMA Platform                                              |
-------------------------------------------------------------------------------

Please enter one of the following commands:
    config: 
        status (filename)
            - Report the current configuration status in total or for the
              specified file
        list_local
            - List available usdotfhwastol configurations on the host machine
            -d 
                - List available usdotfhwastoldev configurations on the host machine
            -c
                - List available usdotfhwastolcandidate configurations on the host machine
        list_remote 
            - List available usdotfhwastol configurations on Dockerhub
            -d
                - List available usdotfhwastoldev configurations on Dockerhub
            -c
                - List available usdotfhwastolcandidate configurations on Dockerhub
        install <tag/image> 
            - Install a configuration identified by <tag> and download 
              dependencies. If <tag> is bare (no :), it is assumed to be a
              usdotfhwastol/carma-config tag.
            -d
                - tag organization is assumed to be usdotfhwastoldev
            -c
                - tag organization is assumed to be usdotfhwastolcandidate
        set <tag/image> 
            - Set the configuration to the version identified by <tag>. If 
              <tag> is bare (no :), it is assumed to be a 
              usdotfhwastol/carma-config tag.
        edit 
            - Open a shell inside the current configuration storage with r/w 
              permissions
        inspect 
            - Open a shell inside the current configuration storage with r/o
              permissions
            -d
                - uses a usdotfhwastoldev/carma-base image
            -c
                - uses a usdotfhwastolcandidate/carma-base image
        reset 
            - Restore a configuration to its default state
    start (all (docker-compose up args)) 
        - Start the CARMA platform's background processes. If all, start 
          everything.
        - Accepts the same flags as "docker-compose up"
    stop (all (docker-compose down args)) 
        - Stop the CARMA platform's foreground processes. If all, stop 
          everything.
        - Accepts the same flags as "docker-compose down"
    exec <optional bash command>
        - Opens a new container with GUI support from the carma-platform or carma-messenger image with the version specified in docker compose. If no version is found it will use latest.
          This container can be used to interact with CARMA using ROS tooling such as Rviz and RQT.
        - If --gui argument is provided it will start docker using nvidia runtime. this flag should be provided as the first flag to work.
        - If the -i <image> argument is provided then the target image will be used instead
        -d
            - uses a usdotfhwastoldev image
        -c
            - uses a usdotfwhastolcandidate image
    ps
        - List all running CARMA docker containers
    attach
        - View STDOUT from all running CARMA foreground processes
    help - Display this information"
HELP
}

carma__config() {
    local cmdname=$1; shift
    if type "carma-config__$cmdname" >/dev/null 2>&1; then
        "carma-config__$cmdname" "$@"
    else
        carma__help
        exit -1
    fi
}

carma() {
    local cmdname=$1; shift
    if type "carma__$cmdname" >/dev/null 2>&1; then
        "carma__$cmdname" "$@"
    else
        carma__help
        exit -1
    fi
}

# if the functions above are sourced into an interactive interpreter, the user can
# just call "carma-config set" or "carma-config reset" with no further code needed.

# if invoked as a script rather than sourced, call function named on argv via the below;
# note that this must be the first operation other than a function definition
# for $_ to successfully distinguish between sourcing and invocation:
[[ $_ != $0 ]] && return

# make sure we actually *did* get passed a valid function name
if declare -f "carma__$1" >/dev/null 2>&1; then
  # invoke that function, passing arguments through
  "carma__$@" # same as "$1" "$2" "$3" ... for full argument list
else
    carma__help
    exit -1
fi
