#!/bin/bash

set -euo pipefail

#############################################################
#### Prometheus Node Exporter Installer                  ####
#### For RHEL/Rocky/Alma, Oracle Linux, Debian & Ubuntu  ####
####                                                     ####
#### Author: Phil Connor                                 ####
#### Contact: contact@mylinux.work                       ####
#### License: MIT                                        ####
#### Version: 1.1                                        ####
####                                                     ####
#### Usage: ./install-node-exporter.sh [OPTIONS]         ####
#############################################################

# Script defaults
INSTALL_DIR="/usr/local/bin"
TEXTFILE_DIR="/var/lib/node_exporter/textfile_collector"
SERVICE_USER="node_exporter"
PORT=9100
COLLECTORS=""
NO_COLLECTORS=""
UPDATE_MODE=false
UNINSTALL_MODE=false
DRY_RUN=false
OPEN_FIREWALL=true
PROMETHEUS_IP=""

# System variables
logfile="/var/log/node-exporter-install.log"
TMPDIR=""

#########################
### Logging Functions ###
#########################
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$logfile"
}

log_error() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $1" | tee -a "$logfile" >&2
}

log_info() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] INFO: $1" | tee -a "$logfile"
}

#########################
### Cleanup Trap      ###
#########################
cleanup() {
    if [[ -n "$TMPDIR" && -d "$TMPDIR" ]]; then
        rm -rf "$TMPDIR"
    fi
}

trap cleanup EXIT

#########################
### Utility Functions ###
#########################
show_help() {
    cat << EOF
Prometheus Node Exporter Installer

USAGE:
    $0 [OPTIONS]

OPTIONS:
    --collectors LIST       Additional collectors to enable (comma-separated)
                            e.g., "systemd,processes,tcpstat"
    --no-collectors LIST    Collectors to disable (comma-separated)
                            e.g., "wifi,infiniband"
    --port PORT             Listen port (default: 9100)
    --textfile-dir DIR      Textfile collector directory
                            (default: /var/lib/node_exporter/textfile_collector)
    --prometheus-ip IP      Restrict firewall rule to this source IP
    --no-firewall           Skip firewall configuration
    --update                Update existing installation
    --uninstall             Remove node_exporter completely
    --dry-run               Show what would be done without doing it
    --help                  Show this help message

EXAMPLES:
    $0
    $0 --collectors "systemd,processes,tcpstat"
    $0 --no-collectors "wifi,infiniband" --port 9200
    $0 --prometheus-ip 10.0.0.5
    $0 --update
    $0 --uninstall
    $0 --dry-run

EOF
}

#########################
### Permission Check  ###
#########################
check_permissions() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root! Login as root, or use sudo."
        exit 1
    fi
}

#########################
### System Detection  ###
#########################
detect_os() {
    if [[ "$(command -v lsb_release)" ]]; then
        OS=$(lsb_release -i | awk '{print $3}' | tr '[:upper:]' '[:lower:]')
        OSVER=$(lsb_release -r | awk '{print $2}' | cut -d. -f1)
    else
        OS=$({ grep PRETTY_NAME /etc/os-release || true; } | sed 's/PRETTY_NAME=//g' | tr -d '="' | awk '{print $1}' | tr '[:upper:]' '[:lower:]')
        OSVER=$({ grep VERSION_ID /etc/os-release || true; } | sed 's/VERSION_ID=//g' | tr -d '"' | cut -d. -f1)
    fi

    log_info "Detected OS: $OS version $OSVER"
}

detect_arch() {
    local machine
    machine=$(uname -m)
    case "$machine" in
        x86_64)  ARCH="amd64" ;;
        aarch64) ARCH="arm64" ;;
        *)
            log_error "Unsupported architecture: $machine"
            exit 1
            ;;
    esac
    log_info "Detected architecture: $ARCH"
}

#########################
### Package Manager   ###
#########################
setup_package_manager() {
    case $OS in
        "ubuntu"|"debian")
            pkgmgr="apt -y"
            ;;
        "red"|"centos"|"oracle"|"rocky"|"almalinux")
            if command -v dnf >/dev/null 2>&1; then
                pkgmgr="dnf -y"
            else
                pkgmgr="yum -y"
            fi
            ;;
        *)
            log_error "Unsupported OS: $OS"
            exit 1
            ;;
    esac
    log_info "Using package manager: $pkgmgr"
}

#########################
### Dependencies      ###
#########################
install_dependencies() {
    for cmd in curl tar; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            log_info "Installing $cmd"
            $pkgmgr install "$cmd"
        fi
    done
}

#########################
### Version Helpers   ###
#########################
get_latest_version() {
    local version
    version=$(curl -s https://api.github.com/repos/prometheus/node_exporter/releases/latest \
        | grep '"tag_name"' | cut -d '"' -f 4 | sed 's/^v//')
    if [[ -z "$version" ]]; then
        log_error "Failed to fetch latest version from GitHub API"
        exit 1
    fi
    echo "$version"
}

get_installed_version() {
    if [[ -x "${INSTALL_DIR}/node_exporter" ]]; then
        "${INSTALL_DIR}/node_exporter" --version 2>&1 | head -1 | awk '{print $3}'
    else
        echo ""
    fi
}

#########################
### User Management   ###
#########################
create_service_user() {
    if ! id "$SERVICE_USER" &>/dev/null; then
        log_info "Creating $SERVICE_USER user"
        useradd --no-create-home --shell /usr/sbin/nologin --system "$SERVICE_USER"
    else
        log_info "User $SERVICE_USER already exists"
    fi
}

#########################
### Download & Install ##
#########################
download_and_install() {
    local version="$1"

    TMPDIR=$(mktemp -d /tmp/node-exporter-install-XXXXXX)
    local tarball="node_exporter-${version}.linux-${ARCH}.tar.gz"
    local url="https://github.com/prometheus/node_exporter/releases/download/v${version}/${tarball}"

    log_info "Downloading node_exporter v${version} for ${ARCH}"
    curl -sL -o "${TMPDIR}/${tarball}" "$url" || {
        log_error "Failed to download ${url}"
        exit 1
    }

    log_info "Extracting archive"
    tar -xzf "${TMPDIR}/${tarball}" -C "$TMPDIR"

    log_info "Installing binary to ${INSTALL_DIR}/node_exporter"
    cp "${TMPDIR}/node_exporter-${version}.linux-${ARCH}/node_exporter" "${INSTALL_DIR}/node_exporter"
    chown root:root "${INSTALL_DIR}/node_exporter"
    chmod 755 "${INSTALL_DIR}/node_exporter"

    # SELinux context for RHEL 8+
    if [[ "$OS" == "red" || "$OS" == "rocky" || "$OS" == "almalinux" || "$OS" == "oracle" ]] && [[ "$OSVER" -ge 8 ]]; then
        restorecon -rv "${INSTALL_DIR}/node_exporter" || true
    fi
}

#########################
### Textfile Directory ##
#########################
create_textfile_dir() {
    log_info "Creating textfile collector directory: ${TEXTFILE_DIR}"
    mkdir -p "$TEXTFILE_DIR"
    chown "$SERVICE_USER":"$SERVICE_USER" "$TEXTFILE_DIR"
}

#########################
### Systemd Service   ###
#########################
build_exec_start() {
    local exec_start="${INSTALL_DIR}/node_exporter"
    exec_start+=" --collector.textfile.directory=${TEXTFILE_DIR}"

    # Custom listen port
    if [[ "$PORT" -ne 9100 ]]; then
        exec_start+=" --web.listen-address=:${PORT}"
    fi

    # Enable additional collectors
    if [[ -n "$COLLECTORS" ]]; then
        IFS=',' read -ra cols <<< "$COLLECTORS"
        for col in "${cols[@]}"; do
            exec_start+=" --collector.${col}"
        done
    fi

    # Disable collectors
    if [[ -n "$NO_COLLECTORS" ]]; then
        IFS=',' read -ra nocols <<< "$NO_COLLECTORS"
        for col in "${nocols[@]}"; do
            exec_start+=" --no-collector.${col}"
        done
    fi

    echo "$exec_start"
}

create_systemd_service() {
    local exec_start
    exec_start=$(build_exec_start)

    log_info "Creating systemd service file"
    cat > /etc/systemd/system/node_exporter.service << EOF
[Unit]
Description=Prometheus Node Exporter
Documentation=https://github.com/prometheus/node_exporter
Wants=network-online.target
After=network-online.target

[Service]
User=${SERVICE_USER}
Group=${SERVICE_USER}
Type=simple
ExecStart=${exec_start}
Restart=always
RestartSec=5s
SyslogIdentifier=node_exporter

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    systemctl enable node_exporter
    systemctl start node_exporter
    log_info "node_exporter service started"
}

#########################
### Firewall Config   ###
#########################
configure_firewall() {
    if [[ "$OPEN_FIREWALL" == "false" ]]; then
        log_info "Skipping firewall configuration (--no-firewall)"
        return
    fi

    # UFW (Debian/Ubuntu)
    if command -v ufw >/dev/null 2>&1 && ufw status | grep -q "active"; then
        log_info "Configuring UFW firewall rule for port ${PORT}"
        if [[ -n "$PROMETHEUS_IP" ]]; then
            ufw allow from "$PROMETHEUS_IP" to any port "$PORT" proto tcp comment "node_exporter" >/dev/null
        else
            ufw allow "$PORT"/tcp comment "node_exporter" >/dev/null
        fi
        log_info "UFW rule added"
        return
    fi

    # firewalld (RHEL/Rocky/Alma/Oracle)
    if command -v firewall-cmd >/dev/null 2>&1 && systemctl is-active --quiet firewalld 2>/dev/null; then
        log_info "Configuring firewalld rule for port ${PORT}"
        if [[ -n "$PROMETHEUS_IP" ]]; then
            firewall-cmd --permanent --new-zone=node_exporter 2>/dev/null || true
            firewall-cmd --permanent --zone=node_exporter --add-source="$PROMETHEUS_IP" 2>/dev/null || true
            firewall-cmd --permanent --zone=node_exporter --add-port="${PORT}/tcp" 2>/dev/null || true
        else
            firewall-cmd --permanent --add-port="${PORT}/tcp" >/dev/null
        fi
        firewall-cmd --reload >/dev/null
        log_info "firewalld rule added"
        return
    fi

    log_info "No active firewall detected, skipping firewall configuration"
}

#########################
### Verify Install    ###
#########################
verify_installation() {
    log_info "Verifying node_exporter service"

    sleep 2

    if ! systemctl is-active --quiet node_exporter; then
        log_error "node_exporter service is not running"
        systemctl status node_exporter --no-pager | tee -a "$logfile"
        exit 1
    fi

    if curl -sf "http://localhost:${PORT}/metrics" >/dev/null 2>&1; then
        log_info "Metrics endpoint responding at http://localhost:${PORT}/metrics"
    else
        log_error "Metrics endpoint not responding on port ${PORT}"
        exit 1
    fi
}

#########################
### Print Summary     ###
#########################
print_summary() {
    local version
    version=$(get_installed_version)

    echo
    echo "=== Node Exporter Installation Summary ==="
    echo "  Version:       ${version}"
    echo "  Binary:        ${INSTALL_DIR}/node_exporter"
    echo "  Service user:  ${SERVICE_USER}"
    echo "  Port:          ${PORT}"
    echo "  Textfile dir:  ${TEXTFILE_DIR}"
    echo "  Metrics URL:   http://localhost:${PORT}/metrics"
    [[ -n "$COLLECTORS" ]]    && echo "  Enabled:       ${COLLECTORS}"
    [[ -n "$NO_COLLECTORS" ]] && echo "  Disabled:      ${NO_COLLECTORS}"
    echo
    echo "  Check logs at: ${logfile}"
    echo
}

#########################
### Install Mode      ###
#########################
do_install() {
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] Would install node_exporter"
        log_info "[DRY RUN]   User: ${SERVICE_USER}"
        log_info "[DRY RUN]   Port: ${PORT}"
        log_info "[DRY RUN]   Textfile dir: ${TEXTFILE_DIR}"
        log_info "[DRY RUN]   Firewall: ${OPEN_FIREWALL}"
        [[ -n "$COLLECTORS" ]]    && log_info "[DRY RUN]   Enable collectors: ${COLLECTORS}"
        [[ -n "$NO_COLLECTORS" ]] && log_info "[DRY RUN]   Disable collectors: ${NO_COLLECTORS}"
        return
    fi

    local version
    version=$(get_latest_version)
    log_info "Latest version: ${version}"

    create_service_user
    download_and_install "$version"
    create_textfile_dir
    create_systemd_service
    configure_firewall
    verify_installation
    print_summary
}

#########################
### Update Mode       ###
#########################
do_update() {
    if [[ ! -x "${INSTALL_DIR}/node_exporter" ]]; then
        log_error "node_exporter is not installed. Run without --update to install."
        exit 1
    fi

    local current latest
    current=$(get_installed_version)
    latest=$(get_latest_version)

    log_info "Installed version: ${current}"
    log_info "Latest version:    ${latest}"

    if [[ "$current" == "$latest" ]]; then
        log_info "Already up to date (v${current}), nothing to do"
        return
    fi

    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] Would update node_exporter from v${current} to v${latest}"
        return
    fi

    log_info "Updating node_exporter from v${current} to v${latest}"

    systemctl stop node_exporter
    download_and_install "$latest"
    systemctl start node_exporter

    verify_installation

    echo
    echo "=== Node Exporter Update Summary ==="
    echo "  Previous version: ${current}"
    echo "  New version:      ${latest}"
    echo "  Metrics URL:      http://localhost:${PORT}/metrics"
    echo
}

#########################
### Uninstall Mode    ###
#########################
do_uninstall() {
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY RUN] Would uninstall node_exporter"
        log_info "[DRY RUN]   Stop and disable service"
        log_info "[DRY RUN]   Remove ${INSTALL_DIR}/node_exporter"
        log_info "[DRY RUN]   Remove /etc/systemd/system/node_exporter.service"
        log_info "[DRY RUN]   Remove ${TEXTFILE_DIR}"
        log_info "[DRY RUN]   Remove user ${SERVICE_USER}"
        log_info "[DRY RUN]   Remove firewall rule for port ${PORT}"
        return
    fi

    log_info "Uninstalling node_exporter"

    # Stop and disable service
    if systemctl is-active --quiet node_exporter 2>/dev/null; then
        systemctl stop node_exporter
        log_info "Stopped node_exporter service"
    fi
    if systemctl is-enabled --quiet node_exporter 2>/dev/null; then
        systemctl disable node_exporter
        log_info "Disabled node_exporter service"
    fi

    # Remove service file
    if [[ -f /etc/systemd/system/node_exporter.service ]]; then
        rm -f /etc/systemd/system/node_exporter.service
        log_info "Removed systemd service file"
    fi

    systemctl daemon-reload

    # Remove binary
    if [[ -f "${INSTALL_DIR}/node_exporter" ]]; then
        rm -f "${INSTALL_DIR}/node_exporter"
        log_info "Removed binary"
    fi

    # Remove textfile directory
    if [[ -d "$TEXTFILE_DIR" ]]; then
        rm -rf "$TEXTFILE_DIR"
        log_info "Removed textfile directory"
    fi

    # Remove user
    if id "$SERVICE_USER" &>/dev/null; then
        userdel "$SERVICE_USER" 2>/dev/null || true
        log_info "Removed user ${SERVICE_USER}"
    fi

    # Remove firewall rules
    if command -v ufw >/dev/null 2>&1 && ufw status | grep -q "active"; then
        ufw delete allow "$PORT"/tcp 2>/dev/null || true
        log_info "Removed UFW rule"
    elif command -v firewall-cmd >/dev/null 2>&1 && systemctl is-active --quiet firewalld 2>/dev/null; then
        firewall-cmd --permanent --remove-port="${PORT}/tcp" 2>/dev/null || true
        firewall-cmd --permanent --delete-zone=node_exporter 2>/dev/null || true
        firewall-cmd --reload >/dev/null 2>/dev/null || true
        log_info "Removed firewalld rule"
    fi

    log_info "node_exporter has been completely removed"
}

#########################
### Parse Arguments   ###
#########################
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --collectors)
                COLLECTORS="$2"
                shift 2
                ;;
            --no-collectors)
                NO_COLLECTORS="$2"
                shift 2
                ;;
            --port)
                PORT="$2"
                shift 2
                ;;
            --textfile-dir)
                TEXTFILE_DIR="$2"
                shift 2
                ;;
            --prometheus-ip)
                PROMETHEUS_IP="$2"
                shift 2
                ;;
            --no-firewall)
                OPEN_FIREWALL=false
                shift
                ;;
            --update)
                UPDATE_MODE=true
                shift
                ;;
            --uninstall)
                UNINSTALL_MODE=true
                shift
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done
}

#########################
### Main              ###
#########################
main() {
    mkdir -p "$(dirname "$logfile")"
    touch "$logfile"

    parse_arguments "$@"

    log_info "Starting node_exporter installer"
    log_info "Command line: $0 $*"

    check_permissions
    detect_os
    detect_arch
    setup_package_manager
    install_dependencies

    if [[ "$UNINSTALL_MODE" == "true" ]]; then
        do_uninstall
    elif [[ "$UPDATE_MODE" == "true" ]]; then
        do_update
    else
        do_install
    fi

    log_info "Done"
}

# Run main function
main "$@"