#!/bin/sh

# ZimaOS Automatic Upgrade Script
# Supports automatic version detection, multiple mirror sources, and automatic cleanup

set -e

# Configuration
GITHUB_MIRROR="https://github.com/IceWhaleTech/ZimaOS/releases/download"
OSS_MIRROR="https://casaos.oss-cn-shanghai.aliyuncs.com/IceWhaleTech/zimaos-rauc/releases/download"
TEMP_DIR="/tmp/zimaos_upgrade"
DOWNLOAD_TIMEOUT=300
MAX_RETRIES=2

# Target version (set by server)
TARGET_VERSION=""

# Version definitions from server
TEST_VERSION="1.5.3"
ALPHA_VERSION="1.5.3"
BETA_VERSION="1.5.3"
STABLE_VERSION="1.5.3"

# Color output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

# Cleanup function
cleanup() {
    log_info "Cleaning up temporary files..."
    if [ -d "$TEMP_DIR" ]; then
        rm -rf "$TEMP_DIR"
    fi

    # Clean up empty directories
    find /tmp -maxdepth 1 -type d -empty -name "zimaos*" -exec rm -rf {} + 2>/dev/null || true
}

# Auto cleanup on exit
trap cleanup EXIT

# Extract version number - use regex to match pure version number
extract_version() {
    local input="$1"
    # Match semantic version: major.minor.patch[-prerelease][+build]
    local version=$(echo "$input" | grep -oE '[0-9]+\.[0-9]+\.[0-9]+(-[a-zA-Z0-9]+[0-9]*)?' | head -n1)
    if [ -n "$version" ]; then
        echo "$version"
    fi
}

# Get current system version
get_current_version() {
    local version=""
    if [ -f /etc/zimaos-release ]; then
        version=$(extract_version "$(cat /etc/zimaos-release)")
    elif [ -f /etc/os-release ]; then
        local os_version=$(grep VERSION_ID /etc/os-release | cut -d= -f2 | tr -d '"')
        version=$(extract_version "$os_version")
    fi

    if [ -n "$version" ]; then
        echo "$version"
    else
        echo "unknown"
    fi
}

# Get user's current channel
get_current_channel() {
    if [ -f /etc/casaos/installer.conf ]; then
        local channel=$(grep "^channel" /etc/casaos/installer.conf | cut -d= -f2 | tr -d ' #')
        if [ -n "$channel" ]; then
            echo "$channel"
        else
            echo "stable"
        fi
    else
        echo "stable"
    fi
}


# Get latest version based on channel
get_latest_version() {
    local channel="$1"

    case "$channel" in
        "test")
            echo "$TEST_VERSION"
            ;;
        "alpha")
            echo "$ALPHA_VERSION"
            ;;
        "beta")
            echo "$BETA_VERSION"
            ;;
        "stable"|*)
            echo "$STABLE_VERSION"
            ;;
    esac
}

# Test download URL availability
test_download_url() {
    local url="$1"
    local timeout=10

    if command -v curl >/dev/null 2>&1; then
        curl -s --connect-timeout $timeout --max-time $timeout -I "$url" | head -n1 | grep -q "200\|302"
    elif command -v wget >/dev/null 2>&1; then
        wget --spider --timeout=$timeout "$url" >/dev/null 2>&1
    else
        return 1
    fi
}

version_ge() {
    # return 0 if $1 >= $2
    [ "$(printf '%s
' "$1" "$2" | sort -V | head -n1)" = "$2" ]
}

# Get available download URLs
get_download_urls() {
    local version="$1"
    local urls=""

    # Package filename
    local package_name
    if [ "${version}" = "1.5.0" ] || version_ge "${version}" "1.5.1"; then
        package_name="zimaos-x86_64-${version}.raucb"
    else
        package_name="zimaos_zimacube-${version}.raucb"
    fi

    # Build URL list with GitHub as default priority
    local test_urls="${GITHUB_MIRROR}/${version}/${package_name}"
    test_urls="$test_urls ${OSS_MIRROR}/${version}/${package_name}"

    # Test URL availability and return first available URL
    for url in $test_urls; do
        # Skip empty URLs
        [ -z "$url" ] && continue
        if test_download_url "$url"; then
            echo "$url"
            return 0
        else
            log_warn "Link unavailable: $url" >&2
        fi
    done

    # No available URLs found
    echo ""
}

# Download and install from URL
download_and_install() {
    local url="$1"
    local filename=$(basename "$url" | cut -d'?' -f1)
    local filepath="$TEMP_DIR/$filename"
    local retries=0

    # Validate URL
    if [ -z "$url" ]; then
        log_error "Empty URL provided to install_from_url"
        return 1
    fi

    # Download file
    while [ $retries -lt $MAX_RETRIES ]; do
        log_info "Attempting download (try $((retries + 1))): $filename"

        if command -v curl >/dev/null 2>&1; then
            if curl -L --progress-bar --connect-timeout 30 --max-time $DOWNLOAD_TIMEOUT -o "$filepath" "$url"; then
                log_success "Download successful: $filename"
                break
            fi
        elif command -v wget >/dev/null 2>&1; then
            if wget --show-progress --timeout=$DOWNLOAD_TIMEOUT -O "$filepath" "$url"; then
                log_success "Download successful: $filename"
                break
            fi
        else
            log_error "Neither wget nor curl is available"
            return 1
        fi

        retries=$((retries + 1))
        [ $retries -lt $MAX_RETRIES ] && sleep 5
    done

    # Check if download succeeded
    if [ $retries -ge $MAX_RETRIES ]; then
        log_error "Download failed after $MAX_RETRIES attempts"
        return 1
    fi

    # Install the downloaded file
    log_info "Installing $filename..."
    if rauc install "$filepath"; then
        return 0
    else
        return 1
    fi
}


# Main upgrade logic
main() {
    log_info "ZimaOS upgrade script starting..."

    # Create temporary directory
    mkdir -p "$TEMP_DIR"

    # Get current version and channel
    local current_version=$(get_current_version)
    local current_channel=$(get_current_channel)
    log_info "Current version: $current_version"
    log_info "Current channel: $current_channel"

    # Determine target version
    local target_version="$TARGET_VERSION"
    if [ -z "$target_version" ]; then
        log_info "Getting latest version for channel: $current_channel"
        target_version=$(get_latest_version "$current_channel")
    fi
    log_info "Target version: $target_version"

    # Check if current version is below 1.2.5 and target is above 1.2.5
    local is_intermediate_upgrade=0
    if [ "$current_version" != "unknown" ] && ! version_ge "$current_version" "1.2.5"; then
        if version_ge "$target_version" "1.2.5"; then
            log_warn "Your current version ($current_version) is below 1.2.5"
            log_warn "Upgrading to 1.2.5 first (required intermediate version)..."
            echo ""

            # Override target version to 1.2.5 and set flag
            target_version="1.2.5"
            is_intermediate_upgrade=1
        fi
    fi

    # Get download URL
    local download_url=$(get_download_urls "$target_version")

    if [ -z "$download_url" ]; then
        log_error "No available download links"
        exit 1
    fi

    # Download and install
    log_info "Starting system upgrade..."
    if download_and_install "$download_url"; then
        log_success "Upgrade completed!"

        # Check if this is an intermediate upgrade to 1.2.5
        if [ $is_intermediate_upgrade -eq 1 ]; then
            echo ""
            echo -e "${YELLOW}========================================${NC}"
            echo -e "${YELLOW}  IMPORTANT: TWO-STEP UPGRADE REQUIRED  ${NC}"
            echo -e "${YELLOW}========================================${NC}"
            echo ""
            log_warn "You have been upgraded to version 1.2.5 (intermediate version)"
            log_warn "After rebooting, please run the upgrade command again to complete"
            log_warn "the upgrade to the target version"
            echo ""
            echo -e "${YELLOW}Next steps:${NC}"
            echo -e "  1. Reboot your system"
            echo -e "  2. Run the upgrade command again after reboot"
            echo ""
            echo -e "${YELLOW}========================================${NC}"
            echo ""

            # Ask if user wants to reboot immediately (only in interactive mode)
            if [ -t 0 ]; then
                echo -n "Reboot now to apply 1.2.5 update? (y/N): "
                read -r answer
                if [ "$answer" = "y" ] || [ "$answer" = "Y" ]; then
                    log_info "Rebooting..."
                    reboot
                fi
            fi
        else
            log_info "System will apply new version after reboot"

            # Ask if user wants to reboot immediately (only in interactive mode)
            if [ -t 0 ]; then
                echo -n "Reboot immediately to apply update? (y/N): "
                read -r answer
                if [ "$answer" = "y" ] || [ "$answer" = "Y" ]; then
                    log_info "Rebooting..."
                    reboot
                else
                    log_info "Please manually reboot system to apply update"
                fi
            else
                log_info "Upgrade completed! Please reboot system to apply update"
            fi
        fi
    else
        log_error "Upgrade failed"
        exit 1
    fi
}


# Execute main function
main