#!/usr/bin/bash

#==========================
# Environment             #
#==========================
export TERM=xterm-256color
export ANSIBLE_HOST_PATTERN_MISMATCH=ignore
export ANSIBLE_PYTHON_INTERPRETER=auto_silent


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

# Ensure unbuffered output (faster logging):
export PYTHONUNBUFFERED=1

# Start with empty tags:
tags=""

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

# Set up logging environment:
provision_log_file_current="$cache_dir/reports/provision.log"
provision_log_file_recent="$cache_dir/reports/provision_recent.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)
emoji_provision="🛠️"
emoji_scheduling="⏱️"
emoji_system="⚙️"
failure_count_file="$cache_dir/state/fail_count"
inhibit=(systemd-inhibit --who=provisioner --why=provisioning)
inventory_url="https://repo.learnlinux.tv/host_configs/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"
target_ansible_version="13.4.0"
unregistered_host_file="/tmp/unregistered_host"
vault_key_file="$config_dir/vault_key"
virtual_env_path="$cache_dir/env"

# Set inventory file:
inventory_file="$cache_dir/state/inventory"


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

ansible_provision_basic() {
    # 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 linux_instance_standard --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
    copy_log_file
    remove_temp_inventory
    exit 0
}

# Provision command (playbook test):
ansible_provision_test() {
    time "$virtual_env_path/bin/ansible-pull" --check --diff \
        --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"
    copy_log_file
    remove_temp_inventory
    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
    copy_log_file
    remove_temp_inventory
    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
    copy_log_file
    capture_provision_date
    remove_temp_inventory
    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"
    set +o pipefail

    # 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
        copy_log_file
        capture_provision_date
        reset_failure_count
        remove_temp_inventory
        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 '%s\n' "$(date +'%Y-%m-%d %H:%M:%S')" > "$sync_date_file"
}

# Capture the date of the last successful full provision:
capture_provision_date() {
    printf '%s\n' "$(date +'%Y-%m-%d %H:%M:%S')" > "$provisioner_date_file"
}

# If the current provision log actually contains useful data,
# then make a copy for later review:
copy_log_file() {
    if [ -f "$provision_log_file_current" ]; then
        if [ "$(stat -c%s "$provision_log_file_current")" -gt 2048 ]; then
            cp "$provision_log_file_current" "$provision_log_file_recent"
        fi
    fi
}

# 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
        --basic|-b)
            collision_check
            root_check
            display_app_info
            prepare_environment
            fetch_inventory
            ansible_provision_basic
            ;;

        --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)
            collision_check
            root_check
            prepare_environment
            fetch_inventory
            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 enabled.\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 "$virtual_env_path" ]; then
                rm -rf "$virtual_env_path"
            fi

            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
            ;;

        --results|-r)
            if [ -f "$provision_log_file_recent" ]; then
                less -R "$provision_log_file_recent"
            else
                printf "Log file not found - has a provision been run yet?\n"
            fi
            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
                    root_check
                    $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
            ;;

        *)
            print_help
            exit 1
            ;;
    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%s is missing.\n" "$private_key_file"
        printf "Please install it and try again.\n"
        exit 1
    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%s is missing.\n" "$vault_key_file"
        printf "Please install it and try again.\n"
        exit 1
    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, empty it it if it does:
    if [ ! -f "$provision_log_file_current" ]; then
        touch "$provision_log_file_current"
    else
        # Empty the current provision log so the file can be re-used:
        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 "ansible==$target_ansible_version"
    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 |dnf5 " > /dev/null; then
        printf "\n${color_red}Collision detected. Exiting to maintain system integrity.${color_white}\n"
        exit 1
    fi

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

    # Loop through each process and check if any of them are running:
    for process in "${bg_processes[@]}"; do
        while pgrep -x "$process" > /dev/null; do
            printf "%s is running.\n" "$process"
            printf "\nThe provision will continue once background processes have finished...\n"
            sleep 10
            clear
        done
    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 --basic   -b\033[0m     Perform a basic provision only\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     Perform a full provision\n"
    printf "  \033[1;36m --results -r\033[0m     View output from the most recent provision\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
}

remove_temp_inventory() {
    if [ -f $unregistered_host_file ]; then
      rm $unregistered_host_file
    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_dir/app.yaml" ]; then
    while IFS= read -r line; do
        [[ "$line" =~ ^[[:space:]]*(#|---|$) ]] && continue
        if [[ "$line" =~ ^([a-zA-Z_][a-zA-Z0-9_]*):[[:space:]]\"(.*)\"$ ]]; then
            export "${BASH_REMATCH[1]}=${BASH_REMATCH[2]}"
        fi
    done < "$config_dir/app.yaml"
elif [ -f "$config_dir/app.cfg" ]; then
    source "$config_dir/app.cfg"
else
    printf "\nNo configuration file was found.\n"
    printf "\nPlease create %s/app.yaml and try again.\n" "$config_dir"
    exit 1
fi

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

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