#!/bin/bash

# Podman detection (set by ensure_docker_or_podman)
PODMAN=""
# Docker alias command (set by ensure_docker_or_podman)
DOCKER=""
# Prefix for all objects in the Docker daemon.
DOCKER_PREFIX=pojde-
# Name of the CA
CA_VOLUME_NAME=${DOCKER_PREFIX}ca
# Container-internal path of the CA
CA_VOLUME_PATH="/opt/pojde/ca"
# Scripts to run to apply/refresh instances. Order matters. "parameters" is run before all others.
SCRIPTS="user apt code-server ttyd novnc jupyter-lab nginx docker pojdectl ssh git modules webwormhole init clean"
# Prefix to display before destructive operations.
CONFIRMATION_PREFIX="This could lead to data loss. Really"
# Global SSH flags for dropbear compatibility
SSH_FLAGS="-oPubkeyAcceptedKeyTypes=+rsa-sha2-512"
# File which contains the current version
VERSION_FILE="${HOME}/.pojdectl_version"
# Endpoint which returns the latest commit hash for pojde
COMMIT_HASH_ENDPOINT="https://api.github.com/repos/pojntfx/pojde/commits/main"

# Shows pojdectl usage information.
print_help() {
    case "$1" in
    "-h" | "--help")
        :
        ;;
    *)
        printf "Unknown command or argument \"${arg}\".\n\n"
        ;;
    esac

    echo "pojdectl is the management tool for pojde.
Global Flags:
[-n]ode <user@host:port>            Remote host to execute on.
                                    If not specified, execute locally.

Modification Commands:
apply <name> <startPort>            Create or upgrade an instance.
    [-f]orce                            Skip confirmation prompts.
    [-u]pgrade                          Pull latest image.
    [-r]ecreate                         Re-create the container.
    [-i]solate                          Block Docker daemon access.
    [-p]rivileged                       Run in privileged mode.
remove [name...]                    Remove instances(s).
    [-f]orce                            Skip confirmation prompts.
    [-c]ustomization                    Remove customizations.
    [-p]references                      Remove preferences.
    [-s]ecurity                         Remove CA.
    [-u]ser data                        Remove user data.
    [-t]ransfer                         Remove transfer data.
    [-d]eb cache                        Remove .deb cache.
    [-a]ll                              Remove everything.
list                                List all instances.

Lifecycle Commands:
start [name...]                     Start instance(s).
stop [name...]                      Stop instance(s).
restart [name...]                   Restart instance(s).

Utility Commands:
logs <name>                                                 Get the logs of an instance.
enter <name>                                                Get a shell in an instance.
forward <name> <local|remote> [lhost:lport:rhost:rport...]  Forward port(s) to or from an instance.

Miscellaneous Commands:
upgrade-pojdectl                    Upgrade this tool.
get-ca-cert [-p]rint                Get the CA cert.
reset-ca [-f]orce                   Reset the CA.

For more information, please visit https://github.com/pojntfx/pojde#Usage."

    exit 0
}

# Asks the user to re-apply their instance.
print_please_reapply() {
    echo 'Please run "pojdectl apply" again to re-initialize.'
}

# Gets a summary of the exposed ports for an instance.
get_port_summary() {
    if [ "$(${DOCKER} inspect -f '{{ .State.Status }}' $1)" = 'running' ]; then
        start_port="$(${DOCKER} inspect -f '{{ (index (index .NetworkSettings.Ports "8000/tcp") 0).HostPort }}' $1)"
        end_port=$((${start_port} + 5))

        echo "${start_port}-${end_port}"
    else
        echo "-"
    fi
}

# Ensures that Docker or Podman are set up
ensure_docker_or_podman() {
    # Detect Podman
    podman_installed=$(if [ "$(command -v podman)" ]; then echo true; else echo false; fi)
    if [ "${podman_installed}" = "true" ]; then
        # Check if podman-remote is installed
        podman_remote_installed=$(if [ "$(command -v podman-remote)" ]; then echo true; else echo false; fi)
        if [ "${podman_remote_installed}" = "true" ]; then
            # Check if podman-remote is running
            status="$(systemctl is-active podman.socket)"
            if [ ! "${status}" = "active" ]; then
                # podman-remote is not running; start it
                sudo systemctl enable --now podman.socket
            fi
        else
            echo "Podman is installed, but podman-remote is missing. Please install podman-remote before continuing."

            exit 1
        fi

        # Alias Docker
        PODMAN="true"
        DOCKER="sudo podman"

        return
    fi

    # Detect Docker
    docker_installed=$(if [ "$(command -v docker)" ]; then echo true; else echo false; fi)
    if [ "${docker_installed}" ]; then
        # Alias Docker to itself
        PODMAN="false"
        DOCKER="docker"

        return
    fi

    # Exit if neither Podman nor Docker are installed
    echo "Neither Docker nor Podman are installed. Please install Docker or Podman before continuing."

    exit 1
}

# Checks if the users has passed -h or --help and if they have done so, display the usage information.
# Also check for the -n flag and set env variables accordingly
run_remotely=false
run_remotely_args=false
for arg in $@; do
    case $arg in
    -h | --help)
        print_help ${arg}
        ;;

    -n)
        run_remotely=true
        ;;

    *)
        if [ "${run_remotely}" = "true" ]; then
            if [ "${run_remotely_args}" = "false" ]; then
                run_remotely_args="${arg}"
            fi
        fi
        ;;
    esac
done

# Run remotely if wanted
if [ "${run_remotely}" = "true" ]; then
    # Split host and port args
    length="$(echo ${run_remotely_args} | tr -cd ':' | wc -c)"
    host=""
    port=""
    IFS=":"
    i=0
    for part in ${run_remotely_args}; do
        if [ "$i" -lt "${length}" ]; then
            if [ "$i" = "0" ]; then
                host="${part}"
            else
                host="${host}:${part}"
            fi
        else
            port="${part}"
        fi

        i=$(($i + 1))
    done
    IFS=" "

    # Exit if neither Podman nor Docker are installed on the remote host
    docker_is_installed=$(if [ -x "$(ssh ${SSH_FLAGS} -p ${port} ${host} command -v docker)" -o -x "$(ssh ${SSH_FLAGS} -p ${port} ${host} command -v podman)" ]; then echo true; else echo false; fi)
    if [ "${docker_is_installed}" = "false" ]; then
        echo "Neither Docker nor Podman are installed on remote host ${host}:${port}. Please install Docker or Podman there before continuing."

        exit 1
    fi

    # Install pojdectl if not already installed on remote host
    is_already_installed=$(if [ -x "$(ssh ${SSH_FLAGS} -p ${port} ${host} command -v pojdectl)" ]; then echo true; else echo false; fi)
    if [ "${is_already_installed}" = "false" ]; then
        scp ${SSH_FLAGS} -P ${port} $(which pojdectl) "$(echo ${host} | sed 's/@\(.*\)/@[\1]/g'):/usr/local/bin"
    fi

    # Strip the `-n` flag from the remote command
    remote_command="$(echo "$@" | sed 's@-n .*:[0-9]\+@@g')"

    # The `forward` command requires advanced steps here
    if [ "$1" != "forward" ]; then
        # Run the command remotely
        ssh ${SSH_FLAGS} -tt -p ${port} ${host} pojdectl $remote_command

        # Don't continue with local execution
        exit 0
    fi
else
    ensure_docker_or_podman
fi

# Handle the main commands
case $1 in
# Create or upgrade an instance.
apply)
    # Read configuration from arguments
    name=""
    start_port=""
    end_port=""
    ssh_port=""
    skip_confirmations=false
    pull_latest_image=false
    recreate_container=false
    isolate=false
    privileged=false
    i=-1
    for arg; do
        i=$((${i} + 1))

        if [ "$i" = "0" ]; then
            continue
        fi

        if [ "${arg}" = "-f" ]; then
            skip_confirmations=true

            continue
        fi

        if [ "${arg}" = "-u" ]; then
            pull_latest_image=true

            continue
        fi

        if [ "${arg}" = "-r" ]; then
            recreate_container=true

            continue
        fi

        if [ "${arg}" = "-i" ]; then
            isolate=true

            continue
        fi

        if [ "${arg}" = "-p" ]; then
            privileged=true

            continue
        fi

        if [ "$i" = "1" ]; then
            name=${arg}

            continue
        fi

        if [ "$i" = "2" ]; then
            start_port=${arg}
            end_port=$((${start_port} + 5))

            continue
        fi

        print_help ${arg}
    done

    # Adjust Docker arguments if host system uses OpenRC
    docker_flags=""
    docker_args=""
    docker_image="pojntfx/pojde:latest"
    if [ ! -n "$(if [ -d /run/systemd/system/ ]; then echo true; fi)" ]; then
        docker_flags="-e POJDE_OPENRC=true"
        docker_args="/sbin/openrc-init"
        docker_image="pojntfx/pojde:latest-openrc"
    fi

    # Pull the latest image
    if [ "${pull_latest_image}" = "true" ]; then
        ${DOCKER} pull ${docker_image}
    fi

    # Enable Docker daemon access
    docker_create_flags=""
    if [ "${isolate}" = "false" ]; then
        docker_create_flags="-v /var/run/docker.sock:/var/run/docker.sock:z"
        # On podman, the podman-remote socket path is different and we have to disable SELinux confinement
        if [ "${PODMAN}" = "true" ]; then
            docker_create_flags="-v /run/podman/podman.sock:/var/run/docker.sock:z --security-opt label=disable"
        fi
    fi

    # Enable privileged mode
    if [ "${privileged}" = "true" ]; then
        docker_create_flags="${docker_create_flags} --privileged"
    fi

    # Re-create the container
    if [ "${recreate_container}" = "true" ]; then
        if [ "${skip_confirmations}" = "true" ]; then
            REPLY='y'
        else
            read -p "${CONFIRMATION_PREFIX} re-create container for ${name} (y/n)? " -n 1 -r
            echo
        fi
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            ${DOCKER} rm -f ${DOCKER_PREFIX}${name}
        fi
    fi

    # Create the container if it doesn't already exist
    if [ ! -n "$(${DOCKER} ps -q -a -f name=${DOCKER_PREFIX}${name})" ]; then
        # Add addtional flags for systemd compatibility (not needed on podman)
        docker_systemd_flags="-e container=oci --tmpfs /tmp:exec --tmpfs /run:exec --tmpfs /run/lock:exec -v /sys/fs/cgroup:/sys/fs/cgroup:ro"
        if [ "${PODMAN}" = "true" ]; then
            # Enable systemd to manipulate it's cgroups configuration
            sudo setsebool -P container_manage_cgroup true

            # Reset the systemd flags, they are not required with podman
            docker_systemd_flags=""
        fi

        ${DOCKER} run ${docker_create_flags} ${docker_systemd_flags} \
            -d \
            --name ${DOCKER_PREFIX}${name} \
            -v ${DOCKER_PREFIX}${name}-preferences:/opt/pojde/preferences:z \
            -v "${CA_VOLUME_NAME}:${CA_VOLUME_PATH}:z" \
            -v ${DOCKER_PREFIX}${name}-home-root:/root:z \
            -v ${DOCKER_PREFIX}${name}-home-user:/home:z \
            -v ${DOCKER_PREFIX}${name}-apt-cache:/var/cache/apt/archives:z \
            -v ${HOME}/Documents/pojde/${name}:/transfer:z \
            -p ${start_port}-${end_port}:8000-8005 \
            --restart always \
            ${docker_flags} \
            ${docker_image} \
            ${docker_args}
    # If the container does already exist, start it
    else
        ${DOCKER} start ${DOCKER_PREFIX}${name}
    fi

    # Ask for parameters
    ${DOCKER} exec -it ${docker_flags} ${DOCKER_PREFIX}${name} bash -c "/opt/pojde/configuration/parameters.sh"

    # Exit if aborted
    if [ "$?" != "0" ]; then
        echo "Apply aborted, exiting."

        exit 1
    fi

    # Run the upgrade hooks of the scripts
    for script in $SCRIPTS; do
        ${DOCKER} exec -it ${docker_flags} ${DOCKER_PREFIX}${name} bash -c ". /opt/pojde/configuration/${script}.sh && upgrade"
    done
    ;;

# Remove instance(s).
remove)
    # Read configuration from arguments
    names=""
    skip_confirmations=false
    remove_customizations=false
    remove_preferences=false
    remove_user_data=false
    remove_transfer_data=false
    remove_deb_cache=false
    i=-1
    for arg; do
        i=$((${i} + 1))

        if [ "$i" = "0" ]; then
            continue
        fi

        if [ "${arg}" = "-f" ]; then
            skip_confirmations=true

            continue
        fi

        if [ "${arg}" = "-c" ]; then
            remove_customizations=true

            continue
        fi

        if [ "${arg}" = "-p" ]; then
            remove_preferences=true

            continue
        fi

        if [ "${arg}" = "-u" ]; then
            remove_user_data=true

            continue
        fi

        if [ "${arg}" = "-t" ]; then
            remove_transfer_data=true

            continue
        fi

        if [ "${arg}" = "-d" ]; then
            remove_deb_cache=true

            continue
        fi

        if [ "${arg}" = "-a" ]; then
            remove_preferences=true
            remove_user_data=true
            remove_deb_cache=true

            continue
        fi

        names="${names} ${arg}"
    done

    # Enable multiple names
    for name in $names; do
        # Remove container
        if [ "${skip_confirmations}" = "true" ]; then
            REPLY='y'
        else
            read -p "${CONFIRMATION_PREFIX} remove container ${name} (y/n)? " -n 1 -r
            echo
        fi
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            ${DOCKER} rm -f ${DOCKER_PREFIX}${name}
        fi

        # Remove customizations
        if [ "${remove_customizations}" = "true" ]; then
            if [ "${skip_confirmations}" = "true" ]; then
                REPLY='y'
            else
                read -p "${CONFIRMATION_PREFIX} remove customizations from ${name} (y/n)? " -n 1 -r
                echo
            fi
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                for script in $SCRIPTS; do
                    ${DOCKER} exec -it ${docker_flags} ${DOCKER_PREFIX}${name} bash -c ". /opt/pojde/configuration/${script}.sh && refresh"
                done

                print_please_reapply
            fi
        fi

        # Remove preferences
        if [ "${remove_preferences}" = "true" ]; then
            if [ "${skip_confirmations}" = "true" ]; then
                REPLY='y'
            else
                read -p "${CONFIRMATION_PREFIX} remove preferences from ${name} (y/n)? " -n 1 -r
                echo
            fi
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                ${DOCKER} volume rm ${DOCKER_PREFIX}${name}-preferences

                print_please_reapply
            fi
        fi

        # Remove user data
        if [ "${remove_user_data}" = "true" ]; then
            if [ "${skip_confirmations}" = "true" ]; then
                REPLY='y'
            else
                read -p "${CONFIRMATION_PREFIX} remove user data from ${name} (y/n)? " -n 1 -r
                echo
            fi
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                ${DOCKER} volume rm ${DOCKER_PREFIX}${name}-home-root ${DOCKER_PREFIX}${name}-home-user

                print_please_reapply
            fi
        fi

        # Remove transfer data
        if [ "${remove_transfer_data}" = "true" ]; then
            if [ "${skip_confirmations}" = "true" ]; then
                REPLY='y'
            else
                read -p "${CONFIRMATION_PREFIX} remove transfer data from ${name} (y/n)? " -n 1 -r
                echo
            fi
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                rm -rf ${HOME}/Documents/pojde/${name}
            fi
        fi

        # Remove .deb cache
        if [ "${remove_deb_cache}" = "true" ]; then
            if [ "${skip_confirmations}" = "true" ]; then
                REPLY='y'
            else
                read -p "${CONFIRMATION_PREFIX} remove .deb cache from ${name} (y/n)? " -n 1 -r
                echo
            fi
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                ${DOCKER} volume rm ${DOCKER_PREFIX}${name}-apt-cache
            fi
        fi
    done
    ;;

# List all instances.
list)
    # Get the current containers' IDs
    container_ids=$(${DOCKER} ps -a --format '{{ .ID }}' -f "name=^${DOCKER_PREFIX}.*")

    # Add a header to the output
    printf "%-30s %-10s %-15s\n" "NAME" "STATUS" "PORTS"

    # Get the name (differs on podman)
    name="/${DOCKER_PREFIX}"
    if [ "${PODMAN}" = "true" ]; then
        name="${DOCKER_PREFIX}"
    fi

    # For each container, show the name, state and exposed ports
    for container_id in $container_ids; do
        printf "%-30s %-10s %-15s\n" "$(${DOCKER} inspect -f "{{ index (split .Name \"${name}\") 1 }}" ${container_id})" "$(${DOCKER} inspect -f "{{ .State.Status }}" ${container_id})" "$(get_port_summary ${container_id})"
    done
    ;;

# Start instance(s).
start)
    i=1
    for arg; do
        if [ "$i" -gt 1 ]; then
            ${DOCKER} start ${DOCKER_PREFIX}${arg}
        fi

        i=$(($i + 1))
    done
    ;;

# Stop instance(s).
stop)
    i=1
    for arg; do
        if [ "$i" -gt 1 ]; then
            ${DOCKER} stop ${DOCKER_PREFIX}${arg}
        fi

        i=$(($i + 1))
    done
    ;;

# Restart instance(s).
restart)
    i=1
    for arg; do
        if [ "$i" -gt 1 ]; then
            ${DOCKER} restart ${DOCKER_PREFIX}${arg}
        fi

        i=$(($i + 1))
    done
    ;;

# Get the logs of an instance.
logs)
    name=$2

    # If container uses systemd, tail from journalctl, else tail OpenRC output
    if [ ! -n "$(${DOCKER} exec -it ${DOCKER_PREFIX}${name} sh -c 'if [ -d /run/systemd/system/ ]; then echo true; fi')" ]; then
        ${DOCKER} exec -it ${DOCKER_PREFIX}${name} tail -f /var/log/messages
    else
        ${DOCKER} exec -it ${DOCKER_PREFIX}${name} journalctl -f
    fi
    ;;

# Get a shell in an instance.
enter)
    name=$2

    ${DOCKER} exec -it ${DOCKER_PREFIX}$2 bash
    ;;

# Forward port(s) from an instance.
forward)
    name=$2
    direction=$3

    # Set the forwarding direction
    directionFlag="-L"
    if [ "${direction}" = "remote" ]; then
        directionFlag="-R"
    fi

    # Strip the `-n` flag if a remote command was specified
    forwardingSpecs="$(echo "$@" | sed 's@-n .*:[0-9]\+@@g')"

    # Construct the forwarding flags
    forwardingFlags=""
    i=1
    for forwardingSpec in $forwardingSpecs; do
        if [ "$i" -gt 3 ]; then
            forwardingFlags="${forwardingFlags} ${directionFlag} ${forwardingSpec}"
        fi

        i=$(($i + 1))
    done

    # Make the port of the SSH server to forward from locally available
    if [ "${run_remotely}" = "true" ]; then
        # Allow accepting fingerprints
        ssh ${SSH_FLAGS} -tt -p ${port} "${host}" "exit 0"

        # Check if the container exists and exit if it does not
        if [ ! -n "$(ssh ${SSH_FLAGS} -p ${port} "${host}" "${DOCKER} ps -q -a -f name=${DOCKER_PREFIX}${name}")" ]; then
            echo "Instance ${DOCKER_PREFIX}${name} could not be found, exiting."

            exit 1
        fi

        # Get the container's SSH port
        remoteSSHPort=$(ssh ${SSH_FLAGS} -p ${port} "${host}" "${DOCKER} inspect -f '{{ (index (index .NetworkSettings.Ports \"8005/tcp\") 0).HostPort }}' ${DOCKER_PREFIX}${name}")

        # Get a random local port
        check="do while"
        while [[ ! -z $check ]]; do
            containerSSHPort=$(((RANDOM % 60000) + 1025))
            check=$(ss -ap | grep $containerSSHPort)
        done

        # Forward the container's SSH port to the random local port
        nohup ssh ${SSH_FLAGS} -p ${port} -f -L 127.0.0.1:${containerSSHPort}:127.0.0.1:${remoteSSHPort} "${host}" -N >/dev/null 2>&1 &
        sshServerPID=$!
    else
        # Check if the container exists and exit if it does not
        if [ ! -n "$(${DOCKER} ps -q -a -f name=${DOCKER_PREFIX}${name})" ]; then
            echo "Instance ${DOCKER_PREFIX}${name} could not be found, exiting."

            exit 1
        fi

        # Get the container's SSH port
        containerSSHPort=$(${DOCKER} inspect -f '{{ (index (index .NetworkSettings.Ports "8005/tcp") 0).HostPort }}' ${DOCKER_PREFIX}${name})
    fi

    # Wait till forwarding is set up
    while ! nc -z 127.0.0.1 ${containerSSHPort}; do
        sleep 0.1
    done

    # Allow accepting fingerprints
    ssh ${SSH_FLAGS} -tt -p ${containerSSHPort} root@127.0.0.1 "exit 0"

    # Forward the ports
    nohup ssh ${SSH_FLAGS} -f -p ${containerSSHPort} ${forwardingFlags} root@127.0.0.1 -N >/dev/null 2>&1 &
    forwardingPID=$!

    # Print the forwarded ports
    echo "${forwardingSpecs}"
    ;;

# Upgrade this tool.
upgrade-pojdectl)
    # Read the following lines into memory first so that replacing the script while it is running doesn't lead to issues
    {
        # Get the current version
        touch "${VERSION_FILE}"
        current_version="$(cat ${VERSION_FILE})"

        # Check the latest version by fetching the commit hash
        latest_version="$(curl -s -H 'Accept: application/vnd.github.VERSION.sha' ${COMMIT_HASH_ENDPOINT})"

        # Upgrade if the versions don't match
        if [ "${current_version}" != "${latest_version}" ]; then
            echo "Upgrading to version \"${latest_version}\" ..."

            # Fetch the latest version from GitHub
            sudo curl -L -o /usr/local/bin/pojdectl https://raw.githubusercontent.com/pojntfx/pojde/main/bin/pojdectl

            # Make it executable
            sudo chmod +x /usr/local/bin/pojdectl

            # Add the new commit hash
            echo "${latest_version}" >"${VERSION_FILE}"

            # Exit to use the new script
            exit
        fi

        echo "No update available. You have version \"${current_version}\", the latest version is \"${latest_version}\"."

        exit
    }

    ;;

# Reset the CA.
reset-ca)
    if [ "${1}" = "-f" ]; then
        REPLY='y'
    else
        read -p "${CONFIRMATION_PREFIX} reset the CA on all containers? This will require re-initialization of all instances. (y/n)" -n 1 -r
        echo
    fi
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        ${DOCKER} volume rm ${CA_VOLUME_NAME}

        echo "Please run \"pojdectl apply\" again for all instances to re-initialize; if you don't do so, you'll loose secure access to them."
    fi
    ;;

# Get the CA.
get-ca-cert)
    if [ "${2}" = "-p" ]; then
        cat "$(${DOCKER} volume inspect -f {{.Mountpoint}} ${CA_VOLUME_NAME})/ca.pem"
    else
        ${DOCKER} exec -it "$(${DOCKER} ps -q -a -f name=${DOCKER_PREFIX} | head -n 1)" ww send "${CA_VOLUME_PATH}/ca.pem"
    fi
    ;;

*)
    print_help $1
    ;;
esac
