#!/bin/bash
# Copyright 2025 Northern.tech AS
#
#    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.

set -e

PROJECT_NAME_ALLOWED_REGEX="[^a-zA-Z0-9_-]"

show_help() {
    cat << EOF
Simple tool to generate Mender Artifacts suitable for the docker-compose module

Usage: $0 [options] [-- [options-for-mender-artifact] ]

    Options: [ -n|--artifact-name -t|--device-type -o|--output_path -a|--architecture -l|--list-architectures -i|--images-dir -m|--manifests-dir -h|--help -p|--project-name]

        --artifact-name      - Artifact name
        --device-type        - Target device type identification (can be given more than once)
        --output-path        - Path to output artifact file. Default: docker-compose-artifact.mender
        --list-architectures - List the available architectures of the images in the manifests
        --architecture       - The architecture to download the images for. See --list-architectures for available architectures.
        --images-dir         - Directory containing container images the composition needs
        --manifests-dir      - Directory containing docker compose manifests
        --project-name       - Name of the docker compose project, must contain only characters from the class $PROJECT_NAME_ALLOWED_REGEX
        --help               - Show help and exit

Anything after a '--' gets passed directly to the mender-artifact tool. See
'mender-artifact write module-image --help' for details. You may want to
at least specify '--software-filesystem' to some keyword describing where
Docker stores images locally, unless it's rootfs (discouraged).
EOF
}

show_help_and_exit_error() {
    show_help
    exit 1
}

check_dependency() {
    if ! which "$1" > /dev/null; then
        echo "The $1 utility is not found but required to generate Artifacts." >&2
        return 1
    fi
}

if ! check_dependency mender-artifact; then
    echo "Please follow the instructions here to install mender-artifact and then try again: https://docs.mender.io/downloads/workstation-tools#mender-artifact" >&2
    exit 1
fi

if ! check_dependency skopeo; then
    echo "Please follow the instructions here to install skopeo and then try again: https://github.com/containers/skopeo" >&2
    exit 1
fi

declare -a device_types
artifact_name=""
output_path="docker-compose-artifact.mender"
project_name=""
passthrough_args=""
version="1.0"
manifests_dir=""
images_dir=""
architecture=""
list_architectures=false

if [ $# -eq 0 ]; then
    show_help_and_exit_error
fi

while test $# -gt 0; do
    case "$1" in
        --project-name | -p)
            if [ -z "$2" ]; then
                show_help_and_exit_error
            fi
            project_name="$2"
            if [[ $project_name =~ $PROJECT_NAME_ALLOWED_REGEX ]]; then
                echo "ERROR: project name must contain only alpha-numerics, _ or -" >&2
                show_help_and_exit_error
            fi
            shift 2
            ;;
        --manifests-dir | -m)
            if [ -z "$2" ]; then
                show_help_and_exit_error
            fi
            manifests_dir="$2"
            shift 2
            ;;
        --images-dir | -i)
            if [ -z "$2" ]; then
                show_help_and_exit_error
            fi
            images_dir="$2"
            shift 2
            ;;
        --device-type | -t)
            if [ -z "$2" ]; then
                show_help_and_exit_error
            fi
            device_types+=("$2")
            shift 2
            ;;
        --artifact-name | -n)
            if [ -z "$2" ]; then
                show_help_and_exit_error
            fi
            artifact_name=$2
            shift 2
            ;;
        --output-path | -o)
            if [ -z "$2" ]; then
                show_help_and_exit_error
            fi
            output_path=$2
            shift 2
            ;;
        --list-architectures | -l)
            list_architectures=true
            shift 1
            ;;
        --architecture | -a)
            if [ -z "$2" ]; then
                show_help_and_exit_error
            fi
            architecture=$2
            shift 2
            ;;
        -h | --help)
            show_help
            exit 0
            ;;
        --)
            shift
            passthrough_args="$@"
            break
            ;;
        -*)
            echo "Error: unsupported option $1" >&2
            show_help_and_exit_error
            ;;
        *)
            shift
            ;;
    esac
done

if [ -z "${manifests_dir}" ]; then
    echo "Directory containing manifests not specified. Aborting." >&2
    show_help_and_exit_error
elif ! [ -d "${manifests_dir}" ]; then
    echo "Manifests directory '${manifests_dir}' doesn't exist. Aborting." >&2
    show_help_and_exit_error
elif [ $(ls -1 "${manifests_dir}" | wc -l) -eq 0 ]; then
    echo "Manifests directory '${manifests_dir}' needs to contain at least one manifest. Aborting." >&2
    show_help_and_exit_error
fi

images=$(sed -n 's/^[[:space:]]*image:[[:space:]]*//p' "${manifests_dir}"/*)
if [ -z "${images}" ]; then
    echo "No images found in manifests. Aborting." >&2
    show_help_and_exit_error
fi

# if true, list and exit - listing architectures only require manifests dir
if [ "$list_architectures" = true ]; then
    for image in $images; do
        archs=$(skopeo inspect --raw docker://"$image" 2> /dev/null | jq -r '([.manifests[]? | select(.platform.os == "linux") | .platform.architecture] | unique | join(","))' 2> /dev/null)
        if [ -z "$archs" ]; then
            # not a multi-arch image
            archs=$(skopeo inspect docker://"$image" 2> /dev/null | jq -r '.Architecture' 2> /dev/null)
        fi
        echo "$image: $archs"
    done
    exit 0
fi

if [ -z "${artifact_name}" ]; then
    echo "Artifact name not specified. Aborting." >&2
    show_help_and_exit_error
fi

if [ -z "${device_types}" ]; then
    echo "Device type not specified. Aborting." >&2
    show_help_and_exit_error
fi

if [ -z "${project_name}" ]; then
    echo "Project name not specified. Aborting." >&2
    show_help_and_exit_error
fi

temp_dir="$(mktemp -d)"
if [[ "${temp_dir}" == "" ]]; then
    echo "Cannot setup temporary directory. Aborting." >&2
    exit 1
fi
function cleanup() {
    rm -rf "$temp_dir"
}
trap cleanup EXIT SIGQUIT SIGTERM

if [ -z "${images_dir}" ]; then
    mkdir $temp_dir/images
    for image in $images; do
        file_name=$(echo "$image" | tr '/:@' '_')
        echo "Downloading image: $image"
        if ! skopeo copy ${architecture:+--override-arch "$architecture"} docker://"$image" docker-archive:"$temp_dir/images/${file_name}.tar":"$image"; then
            echo "ERROR: Failed to download image: $image" >&2
            exit 1
        fi
    done
elif ! [ -d "${images_dir}" ]; then
    echo "Images directory '${images_dir}' doesn't exist. Aborting." >&2
    show_help_and_exit_error
elif [ $(ls -1 "${images_dir}" | wc -l) -eq 0 ]; then
    echo "Images directory '${images_dir}' needs to contain at least one image. Aborting." >&2
    show_help_and_exit_error
else
    cp -r "$images_dir" "${temp_dir}/images"
fi

cp -r "$manifests_dir" "${temp_dir}/manifests"

tar -C "${temp_dir}" -cf "${temp_dir}/manifests.tar" manifests
tar -C "${temp_dir}" -czf "${temp_dir}/images.tar.gz" images
cat << EOF > "${temp_dir}/meta.json"
{"version": "1", "project_name": "$project_name"}
EOF

write_artifact_cmd="mender-artifact write module-image"
write_artifact_cmd+=" --type docker-compose"
write_artifact_cmd+=" --compression none"
write_artifact_cmd+=" --artifact-name ${artifact_name}"
write_artifact_cmd+=" -f ${temp_dir}/images.tar.gz -f ${temp_dir}/manifests.tar"
write_artifact_cmd+=" --meta-data ${temp_dir}/meta.json"
write_artifact_cmd+=" --output-path ${output_path}"

# Just make sure the project name is in the software name and that it's clear
# that it's handled by the docker-compose Update Module.
# We cannot set --software-filesystem because we don't know where docker stores
# images. The user has to do that themselves (and there is a hint about this in
# the --help output).
write_artifact_cmd+=" --software-name mender-docker-compose.${project_name}"

# There can only be one project handled by the docker-compose Update Module so
# every new one installed removes any old installed before.
write_artifact_cmd+=" --clears-provides *mender-docker-compose_*"
write_artifact_cmd+=" --no-default-clears-provides"

for ((i = 0; i < ${#device_types[@]}; i++)); do
    write_artifact_cmd+=" --device-type ${device_types[${i}]}"
done

$write_artifact_cmd $passthrough_args
rc=$?

if [ $rc = 0 ]; then
    echo
    echo "Artifact successfully generated"
    echo
    mender-artifact read "${output_path}"
else
    echo "Failed to generate artifact, see the above output for details."
fi
