#!/bin/bash

################################################
#### Salt Master/Minion Setup Automation    ####
#### Install and configure SaltStack        ####
####                                        ####
#### Author: Phil Connor                    ####
#### Contact: contact@mylinux.work          ####
#### Version: 1.00-030526                   ####
################################################

set -o pipefail

SCRIPT_NAME=$(basename "$0")
readonly SCRIPT_NAME

# Default configuration
readonly DEFAULT_SALT_VERSION="latest"
readonly DEFAULT_FILE_ROOTS="/srv/salt"
readonly DEFAULT_PILLAR_ROOTS="/srv/pillar"
readonly DEFAULT_MASTER_INTERFACE="0.0.0.0"
readonly DEFAULT_MASTER_PORT_PUB=4505
readonly DEFAULT_MASTER_PORT_RET=4506

# Configuration variables (can be overridden by environment)
SALT_VERSION=${SALT_VERSION:-$DEFAULT_SALT_VERSION}
FILE_ROOTS=${FILE_ROOTS:-$DEFAULT_FILE_ROOTS}
PILLAR_ROOTS=${PILLAR_ROOTS:-$DEFAULT_PILLAR_ROOTS}
DEBUG=${DEBUG:-}

# Runtime flags
MODE=""
MASTER_IP=""
MINION_ID=""
AUTO_ACCEPT=false
AUTO_YES=false
PKG_MANAGER=""
OS_FAMILY=""
OS_VERSION=""

handle_error() {
    local exit_code=$1
    local line_number=$2
    echo "Error: $SCRIPT_NAME failed at line $line_number with exit code $exit_code" >&2
    exit "$exit_code"
}

trap 'handle_error $? $LINENO' ERR

debug_echo() {
    if [[ -n "$DEBUG" ]]; then
        echo "[DEBUG] $*" >&2
    fi
}

log_info() {
    echo "[INFO] $*"
}

log_warn() {
    echo "[WARN] $*" >&2
}

log_error() {
    echo "[ERROR] $*" >&2
}

show_help() {
    cat << EOF
Usage: $SCRIPT_NAME [OPTIONS]

Automate Salt master and/or minion installation and configuration.

Supports Ubuntu/Debian and RHEL/AlmaLinux. Adds the Salt Project repository,
installs packages, configures services, creates directory structure, and
opens firewall ports.

OPTIONS:
    --mode master|minion|both   What to install (required)
    --master-ip ADDRESS         Salt master IP or hostname (required for minion/both)
    --minion-id NAME            Custom minion ID (default: system hostname)
    --auto-accept               Enable auto_accept on master (NOT for production)
    --salt-version VERSION      Pin Salt version (default: latest)
    --yes                       Skip confirmation prompts
    --help, -h                  Show this help message

ENVIRONMENT VARIABLES:
    SALT_VERSION        Salt version to install (default: $DEFAULT_SALT_VERSION)
    FILE_ROOTS          Master file_roots path (default: $DEFAULT_FILE_ROOTS)
    PILLAR_ROOTS        Master pillar_roots path (default: $DEFAULT_PILLAR_ROOTS)
    DEBUG               Enable debug output

EXAMPLES:
    # Install salt-master
    sudo $SCRIPT_NAME --mode master --yes

    # Install salt-minion pointing to master
    sudo $SCRIPT_NAME --mode minion --master-ip 10.0.0.1

    # Install both on the same node
    sudo $SCRIPT_NAME --mode both --master-ip localhost --yes

    # Install with custom minion ID
    sudo $SCRIPT_NAME --mode minion --master-ip salt.example.com --minion-id web01

    # Install specific Salt version
    sudo $SCRIPT_NAME --mode master --salt-version 3006 --yes
EOF
}

detect_os() {
    if [[ -f /etc/os-release ]]; then
        # shellcheck disable=SC1091
        source /etc/os-release
        OS_VERSION="$VERSION_ID"
        case "$ID" in
            ubuntu|debian)
                OS_FAMILY="debian"
                PKG_MANAGER="apt"
                ;;
            rhel|centos|rocky|almalinux|ol|fedora)
                OS_FAMILY="rhel"
                if command -v dnf >/dev/null 2>&1; then
                    PKG_MANAGER="dnf"
                else
                    PKG_MANAGER="yum"
                fi
                ;;
            *)
                log_error "Unsupported OS: $ID"
                exit 1
                ;;
        esac
    else
        log_error "Cannot detect OS — /etc/os-release not found"
        exit 1
    fi
    debug_echo "Detected OS: $OS_FAMILY ($PKG_MANAGER) version $OS_VERSION"
}

get_cpu_count() {
    nproc 2>/dev/null || echo 2
}

add_salt_repo_debian() {
    log_info "Adding Salt Project repository (Debian/Ubuntu)..."

    apt-get update -qq
    apt-get install -y -qq curl gnupg2 >/dev/null

    local keyring="/etc/apt/keyrings/salt-archive-keyring.gpg"
    mkdir -p /etc/apt/keyrings
    curl -fsSL "https://repo.saltproject.io/salt/py3/ubuntu/${OS_VERSION}/amd64/SALT-PROJECT-GPG-PUBKEY-2023.gpg" \
        -o "$keyring"

    local repo_url="https://repo.saltproject.io/salt/py3/ubuntu/${OS_VERSION}/amd64"
    if [[ "$SALT_VERSION" != "latest" ]]; then
        repo_url="${repo_url}/${SALT_VERSION}"
    fi
    echo "deb [signed-by=${keyring}] ${repo_url} ${VERSION_CODENAME} main" \
        > /etc/apt/sources.list.d/salt.list

    apt-get update -qq
    log_info "Salt repository added"
}

add_salt_repo_rhel() {
    log_info "Adding Salt Project repository (RHEL)..."

    local major_ver="${OS_VERSION%%.*}"
    local repo_url="https://repo.saltproject.io/salt/py3/redhat/${major_ver}/x86_64"
    if [[ "$SALT_VERSION" != "latest" ]]; then
        repo_url="${repo_url}/${SALT_VERSION}"
    fi

    cat > /etc/yum.repos.d/salt.repo << REPOEOF
[salt]
name=Salt Project for RHEL ${major_ver}
baseurl=${repo_url}
enabled=1
gpgcheck=1
gpgkey=https://repo.saltproject.io/salt/py3/redhat/${major_ver}/x86_64/SALT-PROJECT-GPG-PUBKEY-2023.pub
REPOEOF

    "$PKG_MANAGER" clean expire-cache -q
    log_info "Salt repository added"
}

install_master() {
    log_info "Installing salt-master..."
    case "$PKG_MANAGER" in
        apt)
            apt-get install -y -qq salt-master >/dev/null
            ;;
        dnf|yum)
            "$PKG_MANAGER" install -y -q salt-master
            ;;
    esac
    log_info "salt-master installed"
}

install_minion() {
    log_info "Installing salt-minion..."
    case "$PKG_MANAGER" in
        apt)
            apt-get install -y -qq salt-minion >/dev/null
            ;;
        dnf|yum)
            "$PKG_MANAGER" install -y -q salt-minion
            ;;
    esac
    log_info "salt-minion installed"
}

configure_master() {
    log_info "Configuring salt-master..."

    local worker_threads
    worker_threads=$(get_cpu_count)

    if [[ -f /etc/salt/master ]]; then
        cp /etc/salt/master /etc/salt/master.bak."$(date +%Y%m%d%H%M%S)"
        log_info "Backed up existing /etc/salt/master"
    fi

    cat > /etc/salt/master << MASTEREOF
##### Salt Master Configuration #####
##### Managed by salt-setup.sh  #####

interface: ${DEFAULT_MASTER_INTERFACE}

file_roots:
  base:
    - ${FILE_ROOTS}

pillar_roots:
  base:
    - ${PILLAR_ROOTS}

worker_threads: ${worker_threads}
timeout: 30
state_events: True
presence_events: True
MASTEREOF

    if [[ "$AUTO_ACCEPT" == true ]]; then
        {
            echo ""
            echo "# WARNING: NOT recommended for production"
            echo "auto_accept: True"
        } >> /etc/salt/master
        log_warn "auto_accept enabled — NOT recommended for production"
    else
        {
            echo ""
            echo "auto_accept: False"
        } >> /etc/salt/master
    fi

    log_info "Master configuration written to /etc/salt/master"
}

configure_minion() {
    log_info "Configuring salt-minion..."

    local minion_id
    minion_id="${MINION_ID:-$(hostname -f 2>/dev/null || hostname)}"

    if [[ -f /etc/salt/minion ]]; then
        cp /etc/salt/minion /etc/salt/minion.bak."$(date +%Y%m%d%H%M%S)"
        log_info "Backed up existing /etc/salt/minion"
    fi

    cat > /etc/salt/minion << MINIONEOF
##### Salt Minion Configuration #####
##### Managed by salt-setup.sh  #####

master: ${MASTER_IP}
id: ${minion_id}

# grains:
#   role: webserver
#   environment: production
MINIONEOF

    log_info "Minion configured (id: ${minion_id}, master: ${MASTER_IP})"
}

create_directory_structure() {
    log_info "Creating Salt directory structure..."

    mkdir -p "${FILE_ROOTS}" "${PILLAR_ROOTS}"

    if [[ ! -f "${FILE_ROOTS}/top.sls" ]]; then
        cat > "${FILE_ROOTS}/top.sls" << 'TOPEOF'
base:
  '*':
    []
    # - common
    # - packages
TOPEOF
        log_info "Created ${FILE_ROOTS}/top.sls"
    fi

    if [[ ! -f "${PILLAR_ROOTS}/top.sls" ]]; then
        cat > "${PILLAR_ROOTS}/top.sls" << 'PTOPEOF'
base:
  '*':
    []
    # - common
PTOPEOF
        log_info "Created ${PILLAR_ROOTS}/top.sls"
    fi
}

open_firewall_ports() {
    log_info "Configuring firewall for Salt master ports..."

    if command -v ufw >/dev/null 2>&1; then
        if ufw status | grep -q "Status: active"; then
            ufw allow ${DEFAULT_MASTER_PORT_PUB}/tcp >/dev/null
            ufw allow ${DEFAULT_MASTER_PORT_RET}/tcp >/dev/null
            log_info "Opened ports ${DEFAULT_MASTER_PORT_PUB}/${DEFAULT_MASTER_PORT_RET} in ufw"
        else
            debug_echo "ufw not active — skipping"
        fi
    elif command -v firewall-cmd >/dev/null 2>&1; then
        if firewall-cmd --state >/dev/null 2>&1; then
            firewall-cmd --permanent --add-port=${DEFAULT_MASTER_PORT_PUB}/tcp >/dev/null
            firewall-cmd --permanent --add-port=${DEFAULT_MASTER_PORT_RET}/tcp >/dev/null
            firewall-cmd --reload >/dev/null
            log_info "Opened ports ${DEFAULT_MASTER_PORT_PUB}/${DEFAULT_MASTER_PORT_RET} in firewalld"
        else
            debug_echo "firewalld not running — skipping"
        fi
    else
        log_warn "No supported firewall detected — manually open ports ${DEFAULT_MASTER_PORT_PUB} and ${DEFAULT_MASTER_PORT_RET}"
    fi
}

start_service() {
    local service="$1"
    log_info "Enabling and starting ${service}..."
    systemctl enable "$service" >/dev/null 2>&1
    systemctl restart "$service"
    if systemctl is-active "$service" >/dev/null 2>&1; then
        log_info "${service} is running"
    else
        log_error "${service} failed to start"
        systemctl status "$service" --no-pager
        return 1
    fi
}

show_summary() {
    echo ""
    echo "============================================"
    echo "  Salt Setup Complete"
    echo "============================================"

    if [[ "$MODE" == "master" || "$MODE" == "both" ]]; then
        echo ""
        echo "  Master:"
        echo "    Config:       /etc/salt/master"
        echo "    File roots:   ${FILE_ROOTS}"
        echo "    Pillar roots: ${PILLAR_ROOTS}"
        echo "    Ports:        ${DEFAULT_MASTER_PORT_PUB}, ${DEFAULT_MASTER_PORT_RET}"
        echo ""
        echo "  Master fingerprint:"
        salt-key -F master 2>/dev/null | grep -A1 "master.pub" || echo "    (not yet generated — restart may be needed)"
        echo ""
        echo "  Next steps:"
        echo "    salt-key -L                 # List pending keys"
        echo "    salt-key -a <minion_id>     # Accept a minion key"
        echo "    salt '*' test.ping          # Test connectivity"
    fi

    if [[ "$MODE" == "minion" || "$MODE" == "both" ]]; then
        local minion_id
        minion_id="${MINION_ID:-$(hostname -f 2>/dev/null || hostname)}"
        echo ""
        echo "  Minion:"
        echo "    Config:    /etc/salt/minion"
        echo "    Master:    ${MASTER_IP}"
        echo "    Minion ID: ${minion_id}"
        echo ""
        echo "  Next steps:"
        echo "    salt-call test.ping         # Test master connectivity"
        if [[ "$AUTO_ACCEPT" != true ]]; then
            echo "    (on master) salt-key -a ${minion_id}"
        fi
    fi

    echo ""
    echo "============================================"
}

parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --mode)
                MODE="$2"
                if [[ "$MODE" != "master" && "$MODE" != "minion" && "$MODE" != "both" ]]; then
                    log_error "Mode must be 'master', 'minion', or 'both'"
                    exit 1
                fi
                shift 2
                ;;
            --master-ip)
                MASTER_IP="$2"
                shift 2
                ;;
            --minion-id)
                MINION_ID="$2"
                shift 2
                ;;
            --auto-accept)
                AUTO_ACCEPT=true
                shift
                ;;
            --salt-version)
                SALT_VERSION="$2"
                shift 2
                ;;
            --yes)
                AUTO_YES=true
                shift
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help >&2
                exit 1
                ;;
        esac
    done
}

validate_requirements() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root (use sudo)"
        exit 1
    fi

    if [[ -z "$MODE" ]]; then
        log_error "--mode is required (master, minion, or both)"
        show_help >&2
        exit 1
    fi

    if [[ "$MODE" == "minion" || "$MODE" == "both" ]]; then
        if [[ -z "$MASTER_IP" ]]; then
            log_error "--master-ip is required for minion/both modes"
            exit 1
        fi
    fi

    detect_os
}

main() {
    parse_arguments "$@"
    validate_requirements

    echo "============================================"
    echo "  Salt Setup"
    echo "  Mode: $MODE"
    echo "  OS:   $OS_FAMILY ($PKG_MANAGER)"
    if [[ -n "$MASTER_IP" ]]; then
        echo "  Master: $MASTER_IP"
    fi
    echo "============================================"
    echo ""

    if [[ "$AUTO_YES" != true ]]; then
        echo "Press Enter to continue, or Ctrl+C to abort..."
        read -r
    fi

    case "$OS_FAMILY" in
        debian) add_salt_repo_debian ;;
        rhel)   add_salt_repo_rhel ;;
    esac

    if [[ "$MODE" == "master" || "$MODE" == "both" ]]; then
        install_master
        configure_master
        create_directory_structure
        open_firewall_ports
        start_service salt-master
    fi

    if [[ "$MODE" == "minion" || "$MODE" == "both" ]]; then
        install_minion
        configure_minion
        start_service salt-minion
    fi

    show_summary

    debug_echo "Script completed successfully"
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi