#!/bin/bash

set -euo pipefail

##########################################################################
##  Prometheus Stack Updater                                            ##
##                                                                      ##
##  Updates installed Prometheus ecosystem binaries to latest release   ##
##  from GitHub. Only touches components that are already installed.    ##
##                                                                      ##
##  Supported components:                                               ##
##    prometheus, node_exporter, blackbox_exporter,                     ##
##    alertmanager, mysqld_exporter, promtool, amtool,                  ##
##    loki, promtail, alloy, grafana                                    ##
##                                                                      ##
##  Usage:                                                              ##
##    ./update-prometheus-stack.sh [OPTIONS]                            ##
##                                                                      ##
##  Options:                                                            ##
##    --check           Show what would be updated (no changes)         ##
##    --all             Update all installed components                 ##
##    --prometheus      Update only Prometheus                          ##
##    --node-exporter   Update only node_exporter                       ##
##    --blackbox        Update only blackbox_exporter                   ##
##    --alertmanager    Update only AlertManager                        ##
##    --mysql-exporter  Update only mysqld_exporter                     ##
##    --loki            Update only Loki                                ##
##    --promtail        Update only Promtail                            ##
##    --alloy           Update only Alloy                               ##
##    --grafana         Update only Grafana (via package manager)       ##
##    --force           Update even if already at latest version        ##
##    --arch <arch>     Override architecture (default: auto-detect)    ##
##    --backup-only     Backup configs only (no updates)                ##
##    --help            Show this help message                          ##
##                                                                      ##
##  Author: Phil Connor                                                 ##
##  Contact: pconnor@ara.com                                            ##
##########################################################################

BINDIR="/usr/local/bin"
PROMDIR="/etc/prometheus"
BACKUPDIR="${PROMDIR}/backups"
LOGFILE="/var/log/prometheus-update.log"
TMPDIR_BASE="/tmp/prometheus-update-$$"
CHECK_ONLY=false
BACKUP_ONLY=false
FORCE=false
ARCH=""
UPDATED=0
SKIPPED=0
FAILED=0
COMPONENTS_REQUESTED=()

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

log() {
    local msg
    msg="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    echo -e "$msg" | tee -a "$LOGFILE" 2>/dev/null || echo -e "$msg"
}

log_ok()   { log "${GREEN}✓${NC} $1"; }
log_warn() { log "${YELLOW}⚠${NC} $1"; }
log_err()  { log "${RED}✗${NC} $1" >&2; }
log_info() { log "${CYAN}→${NC} $1"; }

# shellcheck disable=SC2329
cleanup() {
    # shellcheck disable=SC2317
    [[ -d "$TMPDIR_BASE" ]] && rm -rf "$TMPDIR_BASE"
}
trap cleanup EXIT

show_help() {
    sed -n '/^##  Usage:/,/^####/{ /^####/d; s/^##  //; s/^##$//; p }' "$0"
    exit 0
}

detect_arch() {
    if [[ -n "$ARCH" ]]; then
        echo "$ARCH"
        return
    fi
    local machine
    machine=$(uname -m)
    case "$machine" in
        x86_64)  echo "amd64" ;;
        aarch64) echo "arm64" ;;
        armv7l)  echo "armv7" ;;
        armv6l)  echo "armv6" ;;
        *)       echo "amd64" ;;
    esac
}

get_installed_version() {
    local binary="$1"
    local path="${BINDIR}/${binary}"
    if [[ ! -x "$path" ]]; then
        echo "not_installed"
        return
    fi
    case "$binary" in
        prometheus|promtool)
            "$path" --version 2>&1 | head -1 | grep -oP 'version \K[0-9]+\.[0-9]+\.[0-9]+' || echo "unknown"
            ;;
        node_exporter|blackbox_exporter|mysqld_exporter)
            "$path" --version 2>&1 | head -1 | grep -oP 'version \K[0-9]+\.[0-9]+\.[0-9]+' || echo "unknown"
            ;;
        alertmanager|amtool)
            "$path" --version 2>&1 | head -1 | grep -oP 'version \K[0-9]+\.[0-9]+\.[0-9]+' || echo "unknown"
            ;;
        loki|promtail)
            "$path" --version 2>&1 | head -1 | grep -oP 'version \K[0-9]+\.[0-9]+\.[0-9]+' || echo "unknown"
            ;;
        alloy)
            "$path" --version 2>&1 | head -1 | grep -oP '[0-9]+\.[0-9]+\.[0-9]+' || echo "unknown"
            ;;
        *)
            echo "unknown"
            ;;
    esac
}

get_latest_version() {
    local repo="$1"
    local version=""

    case "$repo" in
        prometheus/*)
            local component="${repo#prometheus/}"
            version=$(curl -sf "https://prometheus.io/download/" | \
                grep -oP "${component}-\K[0-9]+\.[0-9]+\.[0-9]+" | head -1 || echo "")
            ;;
        grafana/*)
            version=$(curl -sfL "https://github.com/${repo}/releases/latest" | \
                grep -oP 'releases/tag/v\K[0-9]+\.[0-9]+\.[0-9]+' | head -1 || echo "")
            ;;
    esac

    if [[ -z "$version" ]]; then
        log_err "Failed to query latest version for ${repo}"
        return 1
    fi
    echo "$version"
}

get_download_url() {
    local repo="$1"
    local version="$2"
    local pattern="$3"
    local component="${repo#*/}"

    case "$repo" in
        prometheus/*)
            echo "https://github.com/${repo}/releases/download/v${version}/${component}-${version}.${pattern}"
            ;;
        grafana/*)
            echo "https://github.com/${repo}/releases/download/v${version}/${pattern}"
            ;;
    esac
}

download_and_extract() {
    local url="$1"
    local workdir="$2"
    mkdir -p "$workdir"
    local filename
    filename=$(basename "$url")
    log_info "Downloading ${filename}"
    if ! curl -sfL -o "${workdir}/${filename}" "$url"; then
        log_err "Download failed: ${url}"
        return 1
    fi
    cd "$workdir"
    case "$filename" in
        *.tar.gz|*.tgz)
            tar -xzf "$filename"
            ;;
        *.zip)
            unzip -q "$filename"
            ;;
        *)
            chmod +x "$filename"
            ;;
    esac
}

stop_service() {
    local service="$1"
    if systemctl is-active --quiet "$service" 2>/dev/null; then
        log_info "Stopping ${service}"
        systemctl stop "$service"
        return 0
    fi
    return 1
}

start_service() {
    local service="$1"
    if systemctl is-enabled --quiet "$service" 2>/dev/null; then
        log_info "Starting ${service}"
        systemctl daemon-reload
        systemctl start "$service"
    fi
}

backup_binary() {
    local binary="$1"
    local path="${BINDIR}/${binary}"
    if [[ -f "$path" ]]; then
        local backup
        backup="${path}.backup.$(date +%Y%m%d_%H%M%S)"
        cp "$path" "$backup"
        log_info "Backed up ${path} → ${backup}"
    fi
}

backup_configs() {
    local name="$1"
    local config_files="$2"
    if [[ -z "$config_files" ]]; then
        return 0
    fi
    mkdir -p "$BACKUPDIR"
    local timestamp
    timestamp=$(date +%Y%m%d_%H%M%S)
    for cfg in $config_files; do
        if [[ -f "$cfg" ]]; then
            local filename
            filename=$(basename "$cfg")
            cp "$cfg" "${BACKUPDIR}/${filename}.${timestamp}"
            log_info "Config backed up: ${cfg} → ${BACKUPDIR}/${filename}.${timestamp}"
        fi
    done
}

update_component() {
    local name="$1"
    local repo="$2"
    local service_name="$3"
    local binaries="$4"
    local file_pattern="$5"
    local owner="${6:-prometheus}"
    local config_files="${7:-}"

    local hw
    hw=$(detect_arch)

    local installed
    installed=$(get_installed_version "${binaries%% *}")

    if [[ "$installed" == "not_installed" ]]; then
        return 0
    fi

    local latest
    latest=$(get_latest_version "$repo") || { ((FAILED++)) || true; return 1; }

    echo ""
    log "  ${CYAN}${name}${NC}: installed=${installed}  latest=${latest}"

    if [[ "$installed" == "$latest" ]] && [[ "$FORCE" == "false" ]]; then
        log_ok "Already at latest version"
        ((SKIPPED++)) || true
        return 0
    fi

    if [[ "$CHECK_ONLY" == "true" ]]; then
        if [[ "$installed" != "$latest" ]]; then
            log_warn "Update available: ${installed} → ${latest}"
        fi
        return 0
    fi

    local pattern="${file_pattern//ARCH/${hw}}"
    local url
    url=$(get_download_url "$repo" "$latest" "$pattern")
    if [[ -z "$url" ]]; then
        log_err "Could not find download URL for ${name} (pattern: ${pattern})"
        ((FAILED++)) || true
        return 1
    fi

    local workdir="${TMPDIR_BASE}/${name}"
    download_and_extract "$url" "$workdir" || { ((FAILED++)) || true; return 1; }

    backup_configs "$name" "$config_files"

    local was_running=false
    if stop_service "$service_name"; then
        was_running=true
    fi

    for bin in $binaries; do
        local found
        found=$(find "$workdir" \( -name "$bin" -o -name "${bin}-*" \) -type f 2>/dev/null | head -1)
        if [[ -n "$found" ]]; then
            backup_binary "$bin"
            mv "$found" "${BINDIR}/${bin}"
            chown "${owner}:${owner}" "${BINDIR}/${bin}" 2>/dev/null || \
                chown "${owner}." "${BINDIR}/${bin}" 2>/dev/null || true
            chmod 755 "${BINDIR}/${bin}"
            log_ok "Updated ${bin}"
        else
            log_warn "Binary ${bin} not found in download"
        fi
    done

    if [[ "$was_running" == "true" ]]; then
        start_service "$service_name"
    fi

    local new_ver
    new_ver=$(get_installed_version "${binaries%% *}")
    log_ok "${name} updated: ${installed} → ${new_ver}"
    ((UPDATED++)) || true
}

is_pkg_installed() {
    local pkg="$1"
    if command -v rpm >/dev/null 2>&1; then
        rpm -q "$pkg" >/dev/null 2>&1
    elif command -v dpkg >/dev/null 2>&1; then
        dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"
    else
        return 1
    fi
}

update_alloy() {
    if ! command -v alloy >/dev/null 2>&1 && [[ ! -x "${BINDIR}/alloy" ]]; then
        return 0
    fi

    if is_pkg_installed "alloy"; then
        log_info "Alloy installed via package manager — updating with dnf/apt"
        update_alloy_pkg
    else
        log_info "Alloy installed as standalone binary — updating from GitHub"
        update_component "Alloy" "grafana/alloy" "alloy" "alloy" "alloy-linux-ARCH.zip" "root" "/etc/alloy/config.alloy"
    fi
}

update_alloy_pkg() {
    local alloy_bin="alloy"
    command -v alloy >/dev/null 2>&1 || alloy_bin="${BINDIR}/alloy"

    local installed
    installed=$("$alloy_bin" --version 2>&1 | grep -oP '[0-9]+\.[0-9]+\.[0-9]+' | head -1 || echo "unknown")

    local latest
    latest=$(curl -sfL "https://github.com/grafana/alloy/releases/latest" | \
        grep -oP 'releases/tag/v\K[0-9]+\.[0-9]+\.[0-9]+' | head -1 || echo "")
    if [[ -z "$latest" ]]; then
        log_err "Failed to query latest version for Alloy"
        ((FAILED++)) || true
        return 1
    fi

    echo ""
    log "  ${CYAN}Alloy${NC}: installed=${installed}  latest=${latest}"

    if [[ "$installed" == "$latest" ]] && [[ "$FORCE" == "false" ]]; then
        log_ok "Already at latest version"
        ((SKIPPED++)) || true
        return 0
    fi

    if [[ "$CHECK_ONLY" == "true" ]]; then
        if [[ "$installed" != "$latest" ]]; then
            log_warn "Update available: ${installed} → ${latest}"
        fi
        return 0
    fi

    backup_configs "Alloy" "/etc/alloy/config.alloy"

    if command -v apt-get >/dev/null 2>&1; then
        apt-get -y update && apt-get -y install --only-upgrade alloy
    elif command -v dnf >/dev/null 2>&1; then
        dnf -y upgrade alloy
    elif command -v yum >/dev/null 2>&1; then
        yum -y update alloy
    fi

    systemctl daemon-reload
    systemctl restart alloy

    local new_ver
    new_ver=$("$alloy_bin" --version 2>&1 | grep -oP '[0-9]+\.[0-9]+\.[0-9]+' | head -1 || echo "unknown")
    log_ok "Alloy updated: ${installed} → ${new_ver}"
    ((UPDATED++)) || true
}

update_grafana() {
    if ! command -v grafana-server >/dev/null 2>&1; then
        return 0
    fi

    local installed
    installed=$(grafana-server -v 2>&1 | grep -oP '[0-9]+\.[0-9]+\.[0-9]+' | head -1 || echo "unknown")

    local latest
    latest=$(curl -sfL "https://github.com/grafana/grafana/releases/latest" | \
        grep -oP 'releases/tag/v\K[0-9]+\.[0-9]+\.[0-9]+' | head -1 || echo "")
    if [[ -z "$latest" ]]; then
        log_err "Failed to query latest version for Grafana"
        ((FAILED++)) || true
        return 1
    fi

    echo ""
    log "  ${CYAN}Grafana${NC}: installed=${installed}  latest=${latest}"

    if [[ "$installed" == "$latest" ]] && [[ "$FORCE" == "false" ]]; then
        log_ok "Already at latest version"
        ((SKIPPED++)) || true
        return 0
    fi

    if [[ "$CHECK_ONLY" == "true" ]]; then
        if [[ "$installed" != "$latest" ]]; then
            log_warn "Update available: ${installed} → ${latest}"
        fi
        return 0
    fi

    backup_configs "Grafana" "/etc/grafana/grafana.ini /etc/grafana/ldap.toml"

    log_info "Updating Grafana via package manager"
    if command -v apt-get >/dev/null 2>&1; then
        apt-get -y update && apt-get -y install --only-upgrade grafana
    elif command -v dnf >/dev/null 2>&1; then
        dnf -y upgrade grafana
    elif command -v yum >/dev/null 2>&1; then
        yum -y update grafana
    else
        log_err "No supported package manager found for Grafana update"
        ((FAILED++)) || true
        return 1
    fi

    systemctl daemon-reload
    systemctl restart grafana-server

    local new_ver
    new_ver=$(grafana-server -v 2>&1 | grep -oP '[0-9]+\.[0-9]+\.[0-9]+' | head -1 || echo "unknown")
    log_ok "Grafana updated: ${installed} → ${new_ver}"
    ((UPDATED++)) || true
}

should_update() {
    local component="$1"
    if [[ ${#COMPONENTS_REQUESTED[@]} -eq 0 ]]; then
        return 0
    fi
    for c in "${COMPONENTS_REQUESTED[@]}"; do
        [[ "$c" == "$component" ]] && return 0
    done
    return 1
}

parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --check)       CHECK_ONLY=true; shift ;;
            --backup-only) BACKUP_ONLY=true; shift ;;
            --force)       FORCE=true; shift ;;
            --all)         COMPONENTS_REQUESTED=(); shift ;;
            --prometheus)  COMPONENTS_REQUESTED+=("prometheus"); shift ;;
            --node-exporter) COMPONENTS_REQUESTED+=("node_exporter"); shift ;;
            --blackbox)    COMPONENTS_REQUESTED+=("blackbox"); shift ;;
            --alertmanager) COMPONENTS_REQUESTED+=("alertmanager"); shift ;;
            --mysql-exporter) COMPONENTS_REQUESTED+=("mysql_exporter"); shift ;;
            --loki)        COMPONENTS_REQUESTED+=("loki"); shift ;;
            --promtail)    COMPONENTS_REQUESTED+=("promtail"); shift ;;
            --alloy)       COMPONENTS_REQUESTED+=("alloy"); shift ;;
            --grafana)     COMPONENTS_REQUESTED+=("grafana"); shift ;;
            --arch)        ARCH="$2"; shift 2 ;;
            --help)        show_help ;;
            *)
                log_err "Unknown option: $1"
                show_help
                ;;
        esac
    done
}

main() {
    parse_arguments "$@"

    if [[ $EUID -ne 0 ]]; then
        log_err "This script must be run as root"
        exit 1
    fi

    mkdir -p "$TMPDIR_BASE" "$(dirname "$LOGFILE")"
    touch "$LOGFILE"

    local mode="UPDATE"
    [[ "$CHECK_ONLY" == "true" ]] && mode="CHECK"
    [[ "$BACKUP_ONLY" == "true" ]] && mode="BACKUP"

    echo ""
    echo "=============================================="
    echo "  Prometheus Stack Updater  [${mode}]"
    echo "  $(date '+%Y-%m-%d %H:%M:%S')"
    echo "  Architecture: $(detect_arch)"
    echo "=============================================="

    if [[ "$BACKUP_ONLY" == "true" ]]; then
        local configs=(
            "$PROMDIR/prometheus.yml"
            "$PROMDIR/blackbox.yml"
            "$PROMDIR/alertmanager.yml"
            "/etc/.mysqld_exporter.cnf"
            "/etc/loki/loki-config.yml"
            "/etc/promtail/promtail-config.yml"
            "/etc/alloy/config.alloy"
            "/etc/grafana/grafana.ini"
            "/etc/grafana/ldap.toml"
        )
        local backed_up=0
        mkdir -p "$BACKUPDIR"
        local timestamp
        timestamp=$(date +%Y%m%d_%H%M%S)
        for cfg in "${configs[@]}"; do
            if [[ -f "$cfg" ]]; then
                local filename
                filename=$(basename "$cfg")
                cp "$cfg" "${BACKUPDIR}/${filename}.${timestamp}"
                log_ok "Backed up ${cfg} → ${BACKUPDIR}/${filename}.${timestamp}"
                ((backed_up++))
            fi
        done
        echo ""
        log "Backed up ${backed_up} config file(s) to ${BACKUPDIR}"
        exit 0
    fi

    #                    Name                Repo                                    Service              Binaries                   File Pattern                  Owner              Config Files
    should_update "prometheus"    && update_component "Prometheus"        "prometheus/prometheus"           "prometheus"        "prometheus promtool"          "linux-ARCH.tar.gz"           "prometheus"       "$PROMDIR/prometheus.yml"
    should_update "node_exporter" && update_component "Node Exporter"     "prometheus/node_exporter"       "node_exporter"     "node_exporter"                "linux-ARCH.tar.gz"           "root"             ""
    should_update "blackbox"      && update_component "Blackbox Exporter" "prometheus/blackbox_exporter"   "blackbox_exporter" "blackbox_exporter"            "linux-ARCH.tar.gz"           "prometheus"       "$PROMDIR/blackbox.yml"
    should_update "alertmanager"  && update_component "AlertManager"      "prometheus/alertmanager"        "alertmanager"      "alertmanager amtool"          "linux-ARCH.tar.gz"           "alertmanager"     "$PROMDIR/alertmanager.yml"
    should_update "mysql_exporter" && update_component "MySQL Exporter"   "prometheus/mysqld_exporter"     "mysqld_exporter"   "mysqld_exporter"              "linux-ARCH.tar.gz"           "prometheus"       "/etc/.mysqld_exporter.cnf"
    should_update "loki"          && update_component "Loki"              "grafana/loki"                   "loki"              "loki"                         "loki-linux-ARCH.zip"          "loki"             "/etc/loki/loki-config.yml"
    should_update "promtail"      && update_component "Promtail"          "grafana/loki"                   "promtail"          "promtail"                     "promtail-linux-ARCH.zip"      "promtail"         "/etc/promtail/promtail-config.yml"
    should_update "alloy"         && update_alloy
    should_update "grafana"       && update_grafana

    echo ""
    echo "=============================================="
    echo -e "  Results: ${GREEN}${UPDATED} updated${NC}  ${YELLOW}${SKIPPED} current${NC}  ${RED}${FAILED} failed${NC}"
    echo "=============================================="
    echo ""

    if [[ "$CHECK_ONLY" == "false" ]]; then
        log "Log saved to ${LOGFILE}"
    fi

    [[ $FAILED -gt 0 ]] && exit 1
    exit 0
}

main "$@"