#!/bin/bash

################################################
#### GitLab Upgrade Automation              ####
#### Multi-stop upgrade with air-gapped     ####
#### environment support                    ####
####                                        ####
#### Author: Phil Connor                    ####
#### Contact: contact@mylinux.work          ####
#### Version: 1.01-051326                   ####
################################################

set -o pipefail

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

# Required version stops (as of May 2026)
# Source: https://docs.gitlab.com/update/upgrade_paths/
readonly VERSION_STOPS=(
    "14.0.12"
    "14.3.6"
    "14.9.5"
    "14.10.5"
    "15.0.5"
    "15.4.6"
    "15.11.13"
    "16.0.10"
    "16.3.9"
    "16.7.10"
    "16.11.10"
    "17.3.7"
    "17.5.5"
    "17.8.7"
    "17.11.7"
    "18.0.2"
    "18.2.6"
    "18.5.2"
    "18.8.7"
    "18.11.0"
)

# Default configuration
readonly DEFAULT_PACKAGE_DIR="/var/opt/gitlab/upgrade-packages"
readonly DEFAULT_BACKUP_DIR="/var/opt/gitlab/backups"
readonly DEFAULT_MIN_DISK_GB=5
readonly DEFAULT_MIGRATION_TIMEOUT=3600
readonly DEFAULT_MIGRATION_POLL_INTERVAL=30

# Configuration variables (can be overridden by environment)
PACKAGE_DIR=${PACKAGE_DIR:-$DEFAULT_PACKAGE_DIR}
BACKUP_DIR=${BACKUP_DIR:-$DEFAULT_BACKUP_DIR}
MIN_DISK_GB=${MIN_DISK_GB:-$DEFAULT_MIN_DISK_GB}
MIGRATION_TIMEOUT=${MIGRATION_TIMEOUT:-$DEFAULT_MIGRATION_TIMEOUT}
MIGRATION_POLL_INTERVAL=${MIGRATION_POLL_INTERVAL:-$DEFAULT_MIGRATION_POLL_INTERVAL}
DEBUG=${DEBUG:-}

# Runtime flags
EDITION=""
MODE="upgrade"
AUTO_YES=false
SKIP_BACKUP=false
TARGET_VERSION=""
PKG_MANAGER=""
OS_FAMILY=""
CURRENT_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 multi-stop GitLab upgrades with air-gapped environment support.

Handles required version stops from 14.x through 18.x, background migration
checks, backups, health verification, and supports both online and offline
(air-gapped) installations.

MODES:
    (default)           Perform the upgrade (online or offline)
    --download-only     Download all packages for the upgrade path to PACKAGE_DIR
    --list-stops        Show the upgrade path for your current version and exit

OPTIONS:
    --target VERSION    Target GitLab version (default: latest in stop list)
    --edition ce|ee     GitLab edition (default: auto-detect from installed package)
    --offline           Install from local packages in PACKAGE_DIR instead of repos
    --yes               Skip confirmation prompts between stops
    --skip-backup       Skip the pre-upgrade backup (not recommended)
    --help, -h          Show this help message

ENVIRONMENT VARIABLES:
    PACKAGE_DIR               Directory for downloaded/offline packages (default: $DEFAULT_PACKAGE_DIR)
    BACKUP_DIR                Backup directory (default: $DEFAULT_BACKUP_DIR)
    MIN_DISK_GB               Minimum free disk space in GB (default: $DEFAULT_MIN_DISK_GB)
    MIGRATION_TIMEOUT         Max seconds to wait for background migrations (default: $DEFAULT_MIGRATION_TIMEOUT)
    MIGRATION_POLL_INTERVAL   Seconds between migration status checks (default: $DEFAULT_MIGRATION_POLL_INTERVAL)
    DEBUG                     Enable debug output

EXAMPLES:
    # Show what stops are needed
    $SCRIPT_NAME --list-stops

    # Upgrade to latest (online)
    sudo $SCRIPT_NAME --yes

    # Upgrade to a specific version
    sudo $SCRIPT_NAME --target 17.5.5

    # Download packages for air-gapped transfer
    sudo $SCRIPT_NAME --download-only --target 17.5.5

    # Install from downloaded packages (air-gapped)
    sudo $SCRIPT_NAME --offline --yes

    # Force EE edition
    sudo $SCRIPT_NAME --edition ee --yes
EOF
}

detect_os() {
    if [[ -f /etc/os-release ]]; then
        # shellcheck disable=SC1091
        source /etc/os-release
        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)"
}

detect_edition() {
    if [[ -n "$EDITION" ]]; then
        debug_echo "Edition set by flag: $EDITION"
        return
    fi

    if command -v dpkg >/dev/null 2>&1; then
        if dpkg -l gitlab-ee 2>/dev/null | grep -q "^ii"; then
            EDITION="ee"
        elif dpkg -l gitlab-ce 2>/dev/null | grep -q "^ii"; then
            EDITION="ce"
        fi
    elif command -v rpm >/dev/null 2>&1; then
        if rpm -q gitlab-ee >/dev/null 2>&1; then
            EDITION="ee"
        elif rpm -q gitlab-ce >/dev/null 2>&1; then
            EDITION="ce"
        fi
    fi

    if [[ -z "$EDITION" ]]; then
        log_error "Cannot detect GitLab edition. Use --edition ce|ee"
        exit 1
    fi
    debug_echo "Detected edition: $EDITION"
}

get_current_version() {
    local version_file="/opt/gitlab/embedded/service/gitlab-rails/VERSION"
    if [[ -f "$version_file" ]]; then
        CURRENT_VERSION=$(cat "$version_file")
    else
        log_error "GitLab does not appear to be installed ($version_file not found)"
        exit 1
    fi
    debug_echo "Current version: $CURRENT_VERSION"
}

version_compare() {
    # Returns 0 if $1 < $2, 1 if equal, 2 if $1 > $2
    if [[ "$1" == "$2" ]]; then
        echo 1
        return
    fi

    local IFS=.
    local i
    read -ra ver1 <<< "$1"
    read -ra ver2 <<< "$2"

    for ((i = 0; i < ${#ver1[@]}; i++)); do
        local v1=${ver1[i]:-0}
        local v2=${ver2[i]:-0}
        if ((v1 < v2)); then
            echo 0
            return
        elif ((v1 > v2)); then
            echo 2
            return
        fi
    done
    echo 1
}

get_upgrade_path() {
    local current="$1"
    local target="$2"
    local path=()

    for stop in "${VERSION_STOPS[@]}"; do
        local cmp_current
        cmp_current=$(version_compare "$current" "$stop")
        if [[ "$cmp_current" == "0" ]]; then
            # Current < stop, so this stop is needed
            if [[ -n "$target" ]]; then
                local cmp_target
                cmp_target=$(version_compare "$stop" "$target")
                if [[ "$cmp_target" == "0" || "$cmp_target" == "1" ]]; then
                    path+=("$stop")
                fi
            else
                path+=("$stop")
            fi
        fi
    done

    # If target is specified and not already in the path, add it
    if [[ -n "$target" ]]; then
        local last="${path[*]: -1}"
        if [[ "$last" != "$target" ]]; then
            local in_stops=false
            for stop in "${VERSION_STOPS[@]}"; do
                if [[ "$stop" == "$target" ]]; then
                    in_stops=true
                    break
                fi
            done
            if [[ "$in_stops" == false ]]; then
                path+=("$target")
            fi
        fi
    fi

    echo "${path[@]}"
}

format_package_name() {
    local version="$1"
    if [[ "$OS_FAMILY" == "debian" ]]; then
        echo "gitlab-${EDITION}=${version}-${EDITION}.0"
    else
        echo "gitlab-${EDITION}-${version}-${EDITION}.0"
    fi
}

format_package_filename() {
    local version="$1"
    if [[ "$OS_FAMILY" == "debian" ]]; then
        echo "gitlab-${EDITION}_${version}-${EDITION}.0_amd64.deb"
    else
        echo "gitlab-${EDITION}-${version}-${EDITION}.0.el*.x86_64.rpm"
    fi
}

check_disk_space() {
    local check_path="$1"
    local required_gb="$2"

    local available_kb
    available_kb=$(df --output=avail "$check_path" 2>/dev/null | tail -1 | tr -d ' ')
    local available_gb=$((available_kb / 1024 / 1024))

    if ((available_gb < required_gb)); then
        log_error "Insufficient disk space on $check_path: ${available_gb}GB available, ${required_gb}GB required"
        return 1
    fi
    debug_echo "Disk space on $check_path: ${available_gb}GB available (${required_gb}GB required)"
    return 0
}

check_background_migrations() {
    log_info "Checking background migrations..."

    local pending
    pending=$(sudo gitlab-rails runner -e production \
        'puts Gitlab::Database::BackgroundMigration::BatchedMigration.queued.count' 2>/dev/null) || {
        log_warn "Could not check background migrations via Rails runner"
        return 0
    }

    if [[ "$pending" != "0" ]]; then
        log_warn "$pending background migration(s) still pending"
        log_info "Attempting to finalize background migrations..."
        sudo gitlab-rake db:background_migrations:finalize

        local elapsed=0
        while ((elapsed < MIGRATION_TIMEOUT)); do
            pending=$(sudo gitlab-rails runner -e production \
                'puts Gitlab::Database::BackgroundMigration::BatchedMigration.queued.count' 2>/dev/null) || break
            if [[ "$pending" == "0" ]]; then
                log_info "All background migrations complete"
                return 0
            fi
            log_info "Waiting for $pending migration(s)... (${elapsed}s / ${MIGRATION_TIMEOUT}s)"
            sleep "$MIGRATION_POLL_INTERVAL"
            elapsed=$((elapsed + MIGRATION_POLL_INTERVAL))
        done

        if [[ "$pending" != "0" ]]; then
            log_error "Background migrations did not complete within ${MIGRATION_TIMEOUT}s"
            log_error "Run: sudo gitlab-rake db:background_migrations:status"
            return 1
        fi
    else
        log_info "No pending background migrations"
    fi
    return 0
}

create_backup() {
    if [[ "$SKIP_BACKUP" == true ]]; then
        log_warn "Skipping backup (--skip-backup specified)"
        return 0
    fi

    log_info "Creating GitLab backup..."
    local backup_name
    backup_name="pre-upgrade-$(date +%Y%m%d-%H%M%S)"

    check_disk_space "$BACKUP_DIR" "$MIN_DISK_GB" || return 1

    sudo gitlab-backup create BACKUP="$backup_name" || {
        log_error "Backup failed"
        return 1
    }

    # Back up config files
    log_info "Backing up configuration files..."
    sudo cp /etc/gitlab/gitlab.rb "${BACKUP_DIR}/gitlab.rb.pre-upgrade.bak"
    sudo cp /etc/gitlab/gitlab-secrets.json "${BACKUP_DIR}/gitlab-secrets.json.pre-upgrade.bak"

    log_info "Backup complete: $backup_name"
    return 0
}

verify_health() {
    local version="$1"
    log_info "Verifying GitLab health after upgrade to $version..."

    # Check service status
    if ! sudo gitlab-ctl status >/dev/null 2>&1; then
        log_error "gitlab-ctl status reports unhealthy services"
        sudo gitlab-ctl status
        return 1
    fi
    log_info "All services running"

    # Verify installed version
    local installed
    installed=$(cat /opt/gitlab/embedded/service/gitlab-rails/VERSION 2>/dev/null)
    if [[ "$installed" != "$version" ]]; then
        log_warn "Expected version $version but found $installed"
    else
        log_info "Version confirmed: $installed"
    fi

    # Run health check
    log_info "Running gitlab:check..."
    if ! sudo gitlab-rake gitlab:check SANITIZE=true >/dev/null 2>&1; then
        log_warn "gitlab:check reported issues — review manually"
    else
        log_info "Health check passed"
    fi

    return 0
}

install_package_online() {
    local version="$1"
    local pkg_name
    pkg_name=$(format_package_name "$version")

    log_info "Installing $pkg_name (online)..."

    case "$PKG_MANAGER" in
        apt)
            sudo apt-get update -qq || true
            sudo DEBIAN_FRONTEND=noninteractive apt-get install -y "$pkg_name" || return 1
            ;;
        dnf)
            sudo dnf install -y "$pkg_name" || return 1
            ;;
        yum)
            sudo yum install -y "$pkg_name" || return 1
            ;;
    esac
    return 0
}

install_package_offline() {
    local version="$1"
    local pattern
    pattern=$(format_package_filename "$version")

    # Find the package file
    local pkg_file
    pkg_file=$(find "$PACKAGE_DIR" -name "$pattern" -type f 2>/dev/null | head -1)

    if [[ -z "$pkg_file" ]]; then
        log_error "Package not found in $PACKAGE_DIR matching: $pattern"
        log_error "Run --download-only first to fetch packages"
        return 1
    fi

    log_info "Installing from local package: $pkg_file"

    case "$OS_FAMILY" in
        debian)
            sudo DEBIAN_FRONTEND=noninteractive dpkg -i "$pkg_file" || {
                log_info "Resolving dependencies..."
                sudo apt-get install -f -y || return 1
            }
            ;;
        rhel)
            if [[ "$PKG_MANAGER" == "dnf" ]]; then
                sudo dnf install -y "$pkg_file" || return 1
            else
                sudo yum localinstall -y "$pkg_file" || return 1
            fi
            ;;
    esac
    return 0
}

download_packages() {
    local -a path
    IFS=' ' read -ra path <<< "$(get_upgrade_path "$CURRENT_VERSION" "$TARGET_VERSION")"

    if [[ ${#path[@]} -eq 0 ]]; then
        log_info "No upgrades needed — already at or above target version"
        return 0
    fi

    mkdir -p "$PACKAGE_DIR"

    log_info "Downloading ${#path[@]} package(s) to $PACKAGE_DIR"
    log_info "Upgrade path: ${path[*]}"

    local failed=0
    for version in "${path[@]}"; do
        local pkg_name
        pkg_name=$(format_package_name "$version")
        log_info "Downloading: $pkg_name"

        case "$PKG_MANAGER" in
            apt)
                (cd "$PACKAGE_DIR" && sudo apt-get download "$pkg_name" 2>/dev/null) || {
                    # Try apt-get with download option
                    sudo apt-get install --download-only -y -o Dir::Cache::Archives="$PACKAGE_DIR" "$pkg_name" 2>/dev/null || {
                        log_error "Failed to download $pkg_name"
                        failed=$((failed + 1))
                    }
                }
                ;;
            dnf)
                sudo dnf download --destdir="$PACKAGE_DIR" "$pkg_name" || {
                    log_error "Failed to download $pkg_name"
                    failed=$((failed + 1))
                }
                ;;
            yum)
                sudo yumdownloader --destdir="$PACKAGE_DIR" "$pkg_name" || {
                    log_error "Failed to download $pkg_name"
                    failed=$((failed + 1))
                }
                ;;
        esac
    done

    if ((failed > 0)); then
        log_error "$failed package(s) failed to download"
        return 1
    fi

    log_info "All packages downloaded to $PACKAGE_DIR"
    log_info "Transfer this directory to your air-gapped system, then run:"
    log_info "  PACKAGE_DIR=$PACKAGE_DIR $SCRIPT_NAME --offline"
    return 0
}

list_stops() {
    local -a path
    IFS=' ' read -ra path <<< "$(get_upgrade_path "$CURRENT_VERSION" "$TARGET_VERSION")"

    echo "Current version: $CURRENT_VERSION"
    echo "Edition:         gitlab-$EDITION"
    echo "OS:              $OS_FAMILY ($PKG_MANAGER)"
    if [[ -n "$TARGET_VERSION" ]]; then
        echo "Target version:  $TARGET_VERSION"
    else
        echo "Target version:  latest (${VERSION_STOPS[-1]})"
    fi
    echo ""

    if [[ ${#path[@]} -eq 0 ]]; then
        echo "No upgrades needed — already at or above target version."
        return
    fi

    echo "Required upgrade path (${#path[@]} stops):"
    echo ""
    echo "  $CURRENT_VERSION (current)"
    for stop in "${path[@]}"; do
        echo "    → $stop"
    done
    echo ""
    echo "Estimated time: $((${#path[@]} * 15))-$((${#path[@]} * 30)) minutes (varies by database size)"
}

run_upgrade() {
    local -a path
    IFS=' ' read -ra path <<< "$(get_upgrade_path "$CURRENT_VERSION" "$TARGET_VERSION")"

    if [[ ${#path[@]} -eq 0 ]]; then
        log_info "No upgrades needed — already at or above target version"
        return 0
    fi

    local offline=false
    if [[ "$MODE" == "offline" ]]; then
        offline=true
    fi

    echo "============================================"
    echo "  GitLab Upgrade: $CURRENT_VERSION → ${path[-1]}"
    echo "  Edition: gitlab-$EDITION"
    echo "  Stops: ${#path[@]}"
    echo "  Mode: $(if $offline; then echo "offline"; else echo "online"; fi)"
    echo "============================================"
    echo ""

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

    # Pre-upgrade checks
    check_disk_space "/opt/gitlab" "$MIN_DISK_GB" || exit 1
    check_disk_space "/var/opt/gitlab" "$MIN_DISK_GB" || exit 1
    check_disk_space "/tmp" 1 || exit 1

    # Create backup before starting
    create_backup || exit 1

    # Check initial background migrations
    check_background_migrations || exit 1

    local stop_num=0
    local total=${#path[@]}
    for version in "${path[@]}"; do
        stop_num=$((stop_num + 1))
        echo ""
        echo "============================================"
        echo "  Stop $stop_num/$total: Upgrading to $version"
        echo "============================================"

        # Check background migrations before this stop
        if ((stop_num > 1)); then
            check_background_migrations || {
                log_error "Cannot proceed — background migrations incomplete"
                exit 1
            }
        fi

        # Install the package
        if $offline; then
            install_package_offline "$version" || {
                log_error "Failed to install $version from local packages"
                exit 1
            }
        else
            install_package_online "$version" || {
                log_error "Failed to install $version"
                exit 1
            }
        fi

        # Verify health
        verify_health "$version" || {
            log_warn "Health verification had issues — review before continuing"
        }

        # Check for new background migrations
        log_info "Checking for new background migrations..."
        sudo gitlab-rake db:background_migrations:status 2>/dev/null || true

        if ((stop_num < total)); then
            if [[ "$AUTO_YES" != true ]]; then
                echo ""
                echo "Stop $stop_num/$total complete. Press Enter for next stop, or Ctrl+C to abort..."
                read -r
            else
                log_info "Stop $stop_num/$total complete — proceeding to next stop"
            fi
        fi
    done

    echo ""
    echo "============================================"
    echo "  Upgrade Complete"
    echo "============================================"
    local final_version
    final_version=$(cat /opt/gitlab/embedded/service/gitlab-rails/VERSION 2>/dev/null)
    echo "  Final version: $final_version"
    echo ""
    echo "Post-upgrade tasks:"
    echo "  1. sudo gitlab-rake gitlab:check SANITIZE=true"
    echo "  2. sudo gitlab-rake db:background_migrations:status"
    echo "  3. Verify CI/CD pipelines and runners"
    echo "  4. Check container registry functionality"
    echo "  5. Remove maintenance notification"
}

parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --download-only)
                MODE="download"
                shift
                ;;
            --offline)
                MODE="offline"
                shift
                ;;
            --list-stops)
                MODE="list"
                shift
                ;;
            --target)
                TARGET_VERSION="$2"
                shift 2
                ;;
            --edition)
                EDITION="$2"
                if [[ "$EDITION" != "ce" && "$EDITION" != "ee" ]]; then
                    log_error "Edition must be 'ce' or 'ee'"
                    exit 1
                fi
                shift 2
                ;;
            --yes)
                AUTO_YES=true
                shift
                ;;
            --skip-backup)
                SKIP_BACKUP=true
                shift
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help >&2
                exit 1
                ;;
        esac
    done
}

validate_requirements() {
    # Must run as root for upgrade/download operations
    if [[ "$MODE" != "list" && $EUID -ne 0 ]]; then
        log_error "This script must be run as root (use sudo)"
        exit 1
    fi

    detect_os
    detect_edition
    get_current_version

    # Validate target version format if specified
    if [[ -n "$TARGET_VERSION" ]]; then
        if ! [[ "$TARGET_VERSION" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
            log_error "Invalid target version format: $TARGET_VERSION (expected: X.Y.Z)"
            exit 1
        fi
    fi

    # For offline mode, verify package directory exists
    if [[ "$MODE" == "offline" ]]; then
        if [[ ! -d "$PACKAGE_DIR" ]]; then
            log_error "Package directory not found: $PACKAGE_DIR"
            log_error "Run --download-only first, then transfer packages to this directory"
            exit 1
        fi
    fi
}

main() {
    parse_arguments "$@"

    # Allow --list-stops and --help without root
    if [[ "$MODE" == "list" ]]; then
        detect_os
        detect_edition
        get_current_version
        list_stops
        exit 0
    fi

    validate_requirements

    case "$MODE" in
        download)
            download_packages
            ;;
        upgrade|offline)
            run_upgrade
            ;;
    esac

    debug_echo "Script completed successfully"
}

# Execute main function if script is run directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi