#!/usr/bin/env bash
# By MB
# create a micromamba environment in a apptainer container
# apptainer is preferred, and has more features.

# Usage: micromamba2apptainer.sh <containername> -n <envname> -b -v <version>
function usage() {
    echo "Usage: $0 -c <containername> -n <envname> -v <version>"
    echo "Options:"
    echo "  -c <containername>   Name of the container (default: micromamba)"
    echo "  -n <envname>         Name of the micromamba environment to export (default: base)"
    echo "  -v <version>         Version of the container, only for naming (optional, eg. 20240601 or latest)"
    echo "  -h                   Show this help message"
    echo
    echo "This script requires micromamba and apptainer (1.1.0+) to be installed."
    echo "Example usage:"
    echo " using an existing environment myenv to create a container, with a version tag"
    echo "  $0 -c mycontainer -n myenv -v 202406"
    echo " using an existing environment myenv to create a container"
    echo "  $0 -c mycontainer -n myenv"
    echo " create a new container with a new environment"
    echo "  $0 -c mycontainer"
}

function recipe(){
    cat > Apptainer.recipe <<EOF
Bootstrap: docker
From: mambaorg/micromamba:latest
# This recipe is created by micromamba2apptainer.sh
# University of Vienna, Department of Meteorology and Geophysics
# Author: M. Blaschek
# Date: 06.2025

# rebuild this container with:
# apptainer build --ignore-fakeroot-command --ignore-subuid ${containername}.sif Apptainer.recipe

# Modify the following lines to customize your container
# like: name value
%labels
    maintainer $USER
    version ${version:-latest}
    description "apptainer container with micromamba environment"

# This is the default runscript for the container
# executing the container ./[containername].sif will run this script.
%runscript
    #!/bin/bash
    # is there a command line argument?
    if [ \$# -eq 0 ]; then
        # no
        echo "No arguments provided. Running default shell."
        exec /.singularity.d/actions/shell "\$@"
    else
        # yes
        echo "Running with arguments: \$@"
        exec "\$@"
    fi
EOF
    if [ ! -z "$1" ]; then
        cat >> Apptainer.recipe <<EOF
# this section copies files from the host to the container
# do not change unless you know what you are doing
%files
    # copy the micromamba environment file if it exists
    ${containername}_env_$version.yml /opt/conda/${containername}_env_$version.yml

EOF
    fi
    cat >> Apptainer.recipe <<EOF
# this section is executed during the build of the container
# installing micromamba environment with the packages specified.
%post
    ${micromamba_command}
    micromamba clean --all --yes

# this section defines the environment variables for the container
# used when apptainer shell [containername].sif is executed or the runscript.
%environment
    export PS1="[${containername:-micromamba}]\w\$ "
    export PATH="/opt/conda/bin:\$PATH"
    export LIBRARY="/opt/conda/lib:\$LIBRARY"
    export INCLUDE="/opt/conda/include:\$INCLUDE"
    export CONDA_PREFIX=/opt/conda
    export CONDA_DEFAULT_ENV=base
    export CONDA_EXE=/opt/conda/bin/micromamba
    export CONDA_PROMPT_MODIFIER="(base) "
    export CONDA_PYTHON_EXE=/opt/conda/bin/python
    export CONDA_SHLVL=1
    export CONDA_ROOT=/opt/conda
# end of the Apptainer recipe
# for more information check the documentation at:
# https://apptainer.org/docs/user/main/definition_files.html
EOF
}

if ! command -v micromamba &> /dev/null; then
    echo "Error: micromamba command not found. Please install or load micromamba module and try again."
    exit 1
fi

if ! command -v apptainer &> /dev/null; then
    echo "Error: apptainer command found. Please install or load a module."
    exit 1
else
    apptainer_version=$(apptainer --version | awk '{print $3}')
    required_version="1.1.0"
    if [ "$(printf '%s\n' "$required_version" "$apptainer_version" | sort -V | head -n1)" != "$required_version" ]; then
        echo "Error: apptainer version must be at least $required_version. Current version: $apptainer_version"
        exit 1
    fi
fi

if [ $# -lt 1 ]; then
    usage
    exit 1
fi

while getopts "c:n:v:h" opt; do
    case $opt in
        c)
            containername="$OPTARG"
            ;;
        n)
            envname="$OPTARG"
            ;;
        v)
            version="$OPTARG"
            ;;
        #\?)
        #    echo "Invalid option: -$OPTARG" >&2
        #    exit 1
        #    ;;
        *)
            usage
            exit 1
            ;;
    esac
done

# if no version is specified, use the current date and time as version
if [ -z "$version" ]; then
    version=$(date +%Y%m%d%H%M%S)
fi


# Is there an environment specified?
if [ ! -z "$envname" ]; then
    # yes, check if it exists
    micromamba env list | grep -q "$envname"

    if [ $? -ne 0 ]; then
        echo "Environment '$envname' does not exists. Please choose a different name."
        exit 1
    fi

    # export the environment to a YAML file
    echo "Using existing environment: $envname"
    # using no-build to avoid build numbers in the exported file (when switching from anaconda to conda-forge)
    # or use --from-history to export the environment as it was created, without build numbers.
    micromamba env export --no-build -n "$envname" | head -n -1 | tail -n +2 > "${containername}_env_$version.yml" 
    if [ $? -ne 0 ]; then
        echo "Failed to export environment '$envname'. Please check the error messages."
        exit 1
    fi
    # Create a Apptainer container from the exported environment
    echo "Creating Apptainer container from environment '$envname'..."
    micromamba_command="micromamba -q install -y -n base --file /opt/conda/${containername}_env_$version.yml"
else
    echo "Using micromamba to create a new environment..."
    read -p "Enter Python version (default is 3.12): " python_version
    python_version=${python_version:-3.12}
    echo "If you mistyped, you can edit the Apptainer.recipe file later."
    read -p "Enter all packages to install (space-separated): " packages
    micromamba_command="micromamba -q install -y -n base -c conda-forge python="$python_version" $packages"
fi
echo "Creating Apptainer recipe for container '$containername' with environment '$envname'..."
echo "You can edit that file manually if you want to add more packages or change the base image."
echo "Rebuild: apptainer build --ignore-fakeroot-command --ignore-subuid ${containername}.sif Apptainer.recipe"
echo "------------------------------------------------"
recipe "$envname"
echo "Building Apptainer container..."
# this command should work on almost every system, that has singuarity or apptainer installed.
apptainer build --ignore-fakeroot-command --ignore-subuid "${containername}.sif" Apptainer.recipe
if [ $? -ne 0 ]; then
    echo "Failed to build Apptainer container. Please check the error messages."
    echo
    exit 1
fi
echo "Apptainer container '${containername}.sif' created successfully from environment '$envname'."
echo "------------------------------------------------"
echo "Access non standard directories by setting the SINGULARITY_BIND environment variable."
echo "For example, to include /srvfs, /jetfs or /gpfs (VSC5), run:"
echo "export SINGULARITY_BIND=\"/srvfs,/jetfs,/gpfs\""
echo "You can now run the container with:"
echo "./${containername}.sif"
echo "apptainer run ${containername}.sif"
echo "apptainer shell ${containername}.sif"
echo "To execute a command in the container, use:"
echo "./${containername}.sif <command>"
echo "apptainer exec ${containername}.sif <command>"
echo "------------------------------------------------"
echo "Remember that the micromamba environment is located in /opt/conda."
echo "You can also edit the Apptainer.recipe file to add more packages or change the base image."
echo "To rebuild the container, run:"
echo "apptainer build --ignore-fakeroot-command --ignore-subuid ${containername}.sif Apptainer.recipe"
echo "------------------------------------------------"