#!/usr/bin/bash

#==========================
# Environment             #
#==========================
export TERM=xterm-256color


#==========================
# Variables               #
#==========================
# Set provisioner version:
provisioner_version="3.0.1"

# Start with empty tags:
tags=""

# Set directories:
cache_dir="/opt/provisioner"
config_dir="/etc/provisioner"

# Set up logging environment:
provision_log_file_current="$cache_dir/reports/provision.log"

if [ -f $provision_log_file_current ]; then
    provision_log_file_current_timestamp=$(date -d @"$(stat -c %Y "$provision_log_file_current")" +"%Y-%m-%d_%H-%M-%S")
    provision_log_file_previous="$(echo $provision_log_file_current)_$(echo $provision_log_file_current_timestamp)"
fi

setup_log_file="$cache_dir/reports/setup.log"

# Set configuration values:
app_name_pretty="🤖 \033[1;32mConfig-a-ma-jig 🦾\033[0m"
change_count_file="$cache_dir/state/change_count"
checkout_dir="$cache_dir/repo"
color_red=$(tput setaf 1)
color_white=$(tput setaf 7)
config_file="$config_dir/app.cfg"
emoji_provision="🛠️"
emoji_scheduling="⏱️"
emoji_system="⚙️"
failure_count_file="$cache_dir/state/fail_count"
inhibit="systemd-inhibit --who='provisioner' --why='provisioning';"
inventory_url="http://srv-vps-provisioner-prod.learnlinux.tv/node_assignments"
provisioner_config_version_file="$cache_dir/state/config_version"
provisioner_date_file="$cache_dir/state/last_provision"
self="/usr/bin/provision"
standard_role="linux_instance_generic"
sync_date_file="$cache_dir/state/last_sync"
unregistered_host_file="/tmp/unregistered_host"
unregistered_hostname="unregistered-host"
vault_key_file="$config_dir/vault_key"
virtual_env_path="$cache_dir/env"

# Set inventory file:
if [ -f $unregistered_host_file ]; then
    inventory_file="$unregistered_host_File"
else
    inventory_file="$cache_dir/state/inventory"
fi


#==========================
# Functions               #
#==========================

# Provision command (playbook test):
ansible_provision_test() {
    time "$virtual_env_path/bin/ansible-pull" --check \
        --vault-password-file "$ansible_vault_password_file" --accept-host-key \
        --private-key "$private_key_file" -U "$ansible_repository_url" \
        -C "$repository_branch" -i $inventory_file run.yml
    printf "\nNote: This provision was run in 'check' mode.\n"
    printf "No changes were made to this system.\n"
    exit
}

# Provision command (tags included):
ansible_provision_with_tags() {
    # Save the date of the last provsion attempt for later troubleshooting:
    capture_check_in_date

    time ANSIBLE_HOME="$checkout_dir" \
      stdbuf -oL -eL $inhibit "$virtual_env_path/bin/ansible-pull" \
          --vault-password-file "$ansible_vault_password_file" --accept-host-key \
          --tags $tags --private-key "$private_key_file" -U "$ansible_repository_url" \
          -C "$repository_branch" -i $inventory_file run.yml \
      |& tee -p "$provision_log_file_current"
    record_changes
    exit 0
}

# Provision command (tags NOT included):
ansible_provision_without_tags() {
    # Save the date of the last provsion attempt for later troubleshooting:
    capture_check_in_date

    time ANSIBLE_HOME="$checkout_dir" \
        stdbuf -oL -eL $inhibit "$virtual_env_path/bin/ansible-pull" \
            --vault-password-file "$ansible_vault_password_file" --accept-host-key \
            --private-key "$private_key_file" -U "$ansible_repository_url" \
            -C "$repository_branch" -i "$inventory_file" run.yml \
        |& tee -p "$provision_log_file_current"
    record_changes
    exit 0
}

# Provision command (run via scheduled service):
ansible_provision_scheduled() {
    # Save the date of the last provsion attempt for later troubleshooting:
    capture_check_in_date

    # Fail the pipeline if any command fails; still capture ansible's status explicitly:
    set -o pipefail

    # Run the provision:
    {
        time -p ANSIBLE_HOME="$checkout_dir" $inhibit "$virtual_env_path/bin/ansible-pull" --only-if-changed \
            --vault-password-file "$ansible_vault_password_file" --accept-host-key \
            --private-key "$private_key_file" -U "$ansible_repository_url" \
            -C "$repository_branch" -i $inventory_file run.yml 2>&1
    } | tee -p "$provision_log_file_current"

    # From the pipeline, use PIPESTATUS to get the exit of ansible-pull specifically:
    return_code=${PIPESTATUS[0]}

    # If the job fails, send an alert:
    if [ "$return_code" -ne 0 ]; then
        send_failure_message
        exit 1
    # Otherwise, success:
    else
        record_changes
        reset_failure_count
        exit 0
    fi
}

# Fetch the inventory file from the server:
fetch_inventory() {
    if ! wget -q -O "$inventory_file" "$inventory_url" || \
       ! grep -Fxq "$HOSTNAME" "$inventory_file"; then

        printf "Failed to download host inventory or hostname missing.\nApplying a standard configuration.\n"

        echo "[$standard_role]" > "$unregistered_host_file"
        echo "$HOSTNAME" >> "$unregistered_host_file"

        inventory_file="$unregistered_host_file"
    fi
}

# Capture the date of the current run:
capture_check_in_date() {
    # Capture the most recent check-in:
    printf "$(date +'%Y-%m-%d %H:%M:%S')\n" > $sync_date_file
}

# Display application information:
display_app_info() {
    printf "$app_name_pretty\n\n"
    printf "\033[1;32mSystem Information\033[0m:\n"

    printf "  \033[36mHost\033[0m           $(hostname -f)\n"
    printf "  \033[36mDistribution\033[0m   %s %s\n" \
        "$(lsb_release -si 2>/dev/null || lsb-release -si 2>/dev/null || grep '^NAME=' /etc/os-release 2>/dev/null | cut -d= -f2 | tr -d '"')" \
        "$(lsb_release -sr 2>/dev/null || lsb-release -sr 2>/dev/null || grep '^VERSION_ID=' /etc/os-release 2>/dev/null | cut -d= -f2 | tr -d '"')"
    printf "  \033[36mGroup\033[0m          ${repository_branch^}\n\n"
    printf "\033[1;32mProvisioner Components\033[0m:\n"

    if [ -f $virtual_env_path/bin/ansible ]; then
        printf "  \033[36mAnsible\033[0m        v$($virtual_env_path/bin/ansible --version | head -n 1 | cut -d' ' -f3 | tr -d ']')\n"
    else
        printf "  \033[36mAnsible\033[0m        Not detected\n"
    fi

    if [ -f $provisioner_config_version_file ]; then
        printf "  \033[36mConfig\033[0m         v$(cat $provisioner_config_version_file)\n"
    fi

    printf "  \033[36mProvisioner\033[0m    v$provisioner_version\n\n"
}

# Handle script options:
handle_option() {
    case $1 in
        --disable-timer)
            root_check
            if ! systemctl is-enabled --quiet provision.timer; then
                printf "The provision timer is already disabled.\n"
                exit 1
            else
                systemctl disable provision.timer
                if [ $? -eq 0 ]; then
                    printf "\nThe provision timer was successfully disabled.\n"
                    exit 0
                fi
            fi
            ;;
        --dry-run|-d)
            ansible_provision_test
            exit 0
            ;;
        --enable-timer)
            root_check
            if systemctl is-enabled --quiet provision.timer; then
                printf "The provision timer is already enabled.\n"
                exit 1
            else
                systemctl enable provision.timer
                if [ $? -eq 0 ]; then
                    printf "\nThe provision timer was successfully enbled.\n"
                    exit 0
                fi
            fi
            ;;
        --follow|-f)
            if [ -f $provision_log_file_current ]; then
                tail -f $provision_log_file_current
                exit 0
            else
                printf "Provision log not found.\n"
                printf "Has the provisioner run since the system has been booted?\n"
                exit 1
            fi
            ;;
        --help|-h)
            print_help
            exit 0
            ;;
        --now|-n)
            if [[ -t 0 ]]; then
                collision_check
                root_check
                display_app_info
                prepare_environment
                fetch_inventory
                ansible_provision_without_tags
            else
                collision_check
                prepare_environment
                fetch_inventory
                ansible_provision_scheduled
            fi
            ;;
        --clean|-c)
            root_check
            collision_check

            if [ -d $checkout_dir/pull ]; then
                rm -rf $checkout_dir/pull
            fi

            if [ -d $checkout_dir/tmp ]; then
                rm -rf $checkout_dir/tmp
            fi

            systemctl enable provision.timer
            systemctl restart provision.timer
            exit 0
            ;;
        --role|-r)
            if [ -n "$2" ]; then
                tags=$2
                collision_check
                root_check
                display_app_info
                prepare_environment
                fetch_inventory
                ansible_provision_with_tags
            else
                printf "Error: The --role option was used, but a role was not provided. Use --role <role_name>.\n"
                exit 1
            fi
            ;;
        --status|-s)
            if [ -f $provisioner_date_file ]; then
                printf '%b' "\n$app_name_pretty\n\n"
                printf "\033[36mGroup\033[0m              ${repository_branch^}\n"

                # The $sync_date_file contains the date of the most recent attempt:
                if [ -f $sync_date_file ]; then
                    printf "\033[36mLatest check-in\033[0m %s\n" "   $(cat $sync_date_file)"
                else
                    printf "\033[36mLatest check-in\033[0m   NA\n"
                fi

                printf "\033[36mLatest provision\033[0m %s\n" "  $(cat $provisioner_date_file)"

                if [ ! -f $failure_count_file ]; then
                    if [ -f $change_count_file ]; then
                        recent_changes=$(<$change_count_file)
                        printf "\033[36mRecent changes\033[0m     $recent_changes\n"
                    else
                        printf "\033[36mRecent changes\033[0m     Unknown\n"
                    fi
                else
                    recent_failures=$(<$failure_count_file)
                    printf "${color_red}Failed tasks:      $recent_failures${color_white}\n"
                fi

                # Print when the next provision is scheduled to happen:
                if ps -C ansible-pull > /dev/null 2>&1; then
                    printf "\033[36mNext provision     \033[32mCurrently running...\033[0m\n\n"
                else
                    printf "\033[36mNext provision\033[0m     $(systemctl status provision.timer | grep 'Trigger:' | cut -d ' ' -f 10-11)\n\n"
                fi
            else
                printf "\nIt doesn't appear that a provision has been run since this instance was last started.\n"

                read -p "Do you want to provision it now? (y/n) " resp
                if [ $? -eq 0 ] && [[ $resp =~ ^[yY]$ ]]; then
                    $self
                fi
            fi
            exit 0
            ;;
        --start-timer)
            if systemctl is-active --quiet provision.timer; then
                printf "The provision timer is already enabled.\n"
                exit 1
            else
                systemctl start provision.timer
                if [ $? -eq 0 ]; then
                    printf "The provision timer was successfully started.\n"
                    exit 0
                fi
            fi
            ;;
        --stop-timer)
            if ! systemctl is-active --quiet provision.timer; then
                printf "The provision timer is already stopped.\n"
                exit 1
            else
                systemctl stop provision.timer
                if [ $? -eq 0 ]; then
                    printf "The provision timer was successfully stopped.\n"
                    exit 0
                fi
            fi
            ;;
        --version|-v)
            display_app_info
            exit 0
            ;;
        *)
            if [[ "$1" != "--role" ]]; then
                print_help
                exit 1
            fi
            ;;
    esac
}

# Prepare required app files and settings:
prepare_environment() {
    # The program can't continue without the private key
    if [ ! -f $private_key_file ]; then
        printf "\n$private_key_file is missing.\n"
        printf "Please install it and try again.\n"
    else
        # Ensure permissions for the private key file are correct:
        chown root:root $private_key_file
        chmod 600 $private_key_file
    fi

    # The program can't continue without the vault key:
    if [ ! -f $vault_key_file ]; then
        printf "\n$vault_key_file is missing.\n"
        printf "Please install it and try again.\n"
    else
        chown root:root $vault_key_file
        chmod 600 $vault_key_file
    fi

    # Ensure all required directories exist
    for dir in "$cache_dir" "$cache_dir/reports" "$cache_dir/state"; do
        if [ ! -d "$dir" ]; then
            if mkdir -p "$dir"; then
                printf "\nCreated directory: %s\n" "$dir"
            else
                printf "\nFailed to create directory: %s\nCheck your installation.\n" "$dir"
            fi
        fi
    done

    # Create temporary log if it doesn't exist, rotate it if it does:
    if [ ! -f $provision_log_file_current ]; then
        touch $provision_log_file_current
    else
        if ! grep -q '"changed": false' $provision_log_file_current; then
            cp $provision_log_file_current $provision_log_file_previous
        fi

        truncate -s 0 $provision_log_file_current
    fi

    # Remove older log files that are older than 14 days:
    find $cache_dir/reports -type f -name "*.log*" -mtime +14 -exec rm -f {} \;

    # Create virtual environment
    if [ ! -d $virtual_env_path ]; then
        python3 -m venv $virtual_env_path
    fi

    # Install Ansible and dependencies into the virtual environment:
    if [ ! -f $virtual_env_path/bin/ansible-pull ]; then
        $virtual_env_path/bin/pip install --disable-pip-version-check --upgrade pip
        $virtual_env_path/bin/pip install --disable-pip-version-check --upgrade ansible cryptography jinja2
    fi
}

# Check for app collisions to ensure the provisioner doesn't run against itself:
collision_check() {
    # Check if the provisioner or a package manager is running:
    if pgrep -f "'apt '|ansible-pull|'dnf '" > /dev/null; then
        printf "\n${color_red}Collision detected. Exiting to maintain system integrity.\n"
        exit 1
    fi

    # Wait for background processes to finish if if any are running:
    bg_processes=("apt" "dnf" "packagekitd")

    # Loop through each process and check if any of them are running:
    for process in "${bg_processes[@]}"; do
        if pgrep -x "$process" > /dev/null; then
            printf "$process is running.\n"
            printf "\nThe provision will continue once background processes have finished...\n"
            sleep 10
            clear
        fi
    done
}

# Display help menu:
print_help() {
    printf '%b' "\n$app_name_pretty\n\n"

    printf "System provisioning tool with automatic scheduling\n\n"
    printf 'Usage:\033[3m provision <option>\033[0m\n\n'
    printf '%b' "\033[1;32m$emoji_provision Provision options:\033[0m\n"
    printf "  \033[1;36m --dry-run -d\033[0m     Run a provision, without making any actual changes\n"
    printf "  \033[1;36m --follow  -f\033[0m     Follow the output of an in-progress provision\n"
    printf "  \033[1;36m --now     -n\033[0m     Provision the system immediately\n"
    printf "  \033[1;36m --role    -r\033[0m     Limit the provision to a particular role\n"
    printf "  \033[1;36m --status  -s\033[0m     Show current provision state\n\n"

    printf '%b' "\033[1;32m$emoji_system System options:\033[0m\n"
    printf "  \033[1;36m --clean   -c\033[0m     Clear the cache\n"
    printf "  \033[1;36m --version -v\033[0m     Show app and provisioner versions\n\n"

    printf '%b' "\033[1;32m$emoji_scheduling Scheduling options:\033[0m\n"
    printf "  \033[1;36m --start-timer\033[0m    Start the provision timer\n"
    printf "  \033[1;36m --stop-timer\033[0m     Stop the provision timer\n"
    printf "  \033[1;36m --disable-timer\033[0m  Disable the provision timer for maintenance\n"
    printf "  \033[1;36m --enable-timer\033[0m   Enable the provision timer\n\n"

    printf "   No option (or\033[1;36m --help\033[0m,\033[1;36m -h\033[0m): Display this help menu\n\n"
    exit 0
}

# Retain the number of changes from the previous run:
record_changes() {
    if [ -f $provision_log_file_current ]; then
        # Capture the number of changes during the previous run:
        changed=$(grep -Eo 'changed=[0-9]+' $provision_log_file_current | tail -1 | cut -d= -f2)
        failed=$(grep -Eo 'failed=[0-9]+' $provision_log_file_current | tail -1 | cut -d= -f2)

        # Save the number of changes into a file:
        if [[ -n "$changed" && "$changed" =~ ^[0-9]+$ ]]; then
            echo $changed > $change_count_file
        fi

        # Save the number of failures, if any:
        if [[ -n "$failed" && "$failed" =~ ^[0-9]+$ ]]; then
            if [[ "$failed" -gt 0 ]]; then
                echo $failed > $failure_count_file
            fi
        fi
    fi
}

reset_failure_count() {
    # Reset failure count
    if [ -f $failure_count_file ]; then
        rm $failure_count_file
    fi
}

# Check for root privileges:
root_check() {
    # Note: Only checks for root privileges while run within interactive sessions
    if [ -t 1 ]; then
        if [ $(id -u) -ne 0 ]; then
            printf "\nRun this script as root or using sudo!\n"
            exit 1
        fi
    fi
}

# Send failure message when provisions don't finish:
send_failure_message() {
    pushover_message="⛈️ The Provisioner failed to finish on $HOSTNAME"

    curl -s \
        --form-string "token=$pushover_provisioner_app_token" \
        --form-string "user=$pushover_user_key" \
        --form-string "message=$pushover_message" \
        --form-string "device=$pushover_provisioner_device" https://api.pushover.net/1/messages.json
}


#==========================
# Main Execution          #
#==========================

# Load configuration values:
if [ -f $config_file ]; then
    source $config_file
else
    printf "\n$config_file was not found.\n"
    printf "\nPlease create it at $config_file and try again.\n"
    exit 1
fi

# Process script arguments:
if [ $# -eq 0 ]; then
    print_help
fi

if [ $# -gt 0 ]; then
    handle_option "$1" "$2"
    shift
fi
