#!/bin/bash

####################################################################
#### Code-Server Update Script                                  ####
####                                                            ####
#### Updates code-server to the latest version. Can be run      ####
#### manually or via cron for automatic updates.                ####
####                                                            ####
#### Supported: RHEL/Rocky/Alma 8+, Oracle Linux 8+,            ####
####            Debian 11+, Ubuntu 20.04+                       ####
####                                                            ####
#### Author:  Phil Connor                                       ####
#### Contact: contact@mylinux.work                              ####
#### Website: https://mylinux.work                              ####
#### License: MIT                                               ####
#### Version: 2.0                                               ####
####                                                            ####
#### Usage:                                                     ####
####   sudo ./update-code-server.sh                             ####
####   sudo ./update-code-server.sh --check                     ####
####   sudo ./update-code-server.sh --force                     ####
####   sudo ./update-code-server.sh --install-cron              ####
####   sudo ./update-code-server.sh --help                      ####
####################################################################

set -euo pipefail

# ============================================================================
# USER CONFIGURATION -- edit these or override with command-line options
# ============================================================================

SERVDIR="/usr/local/code-server"      # code-server install directory

# ============================================================================
# INTERNAL VARIABLES
# ============================================================================

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'
VERSION="2.0"
ARCH=""
CHECK_ONLY=0
FORCE=0
INSTALL_CRON=0
CRON_DAYS=25

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

log()      { echo -e "${GREEN}[+]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[!]${NC} $1"; }
log_err()  { echo -e "${RED}[x]${NC} $1" >&2; }
log_info() { echo -e "${CYAN}[>]${NC} $1"; }

show_help() {
    cat <<EOF
Code-Server Update Script v${VERSION}

Usage: sudo $0 [OPTIONS]

Checks for a newer version of code-server and updates if available.

OPTIONS:
    --check             Check for updates without installing (dry run)
    --force             Update even if already on the latest version
    --install-cron      Install this script to /usr/local/bin and set up a
                        cron job for automatic updates (default: every 25 days)
    --cron-days DAYS    Days between automatic updates (default: 25)
    --installdir DIR    code-server install directory (default: /usr/local/code-server)
    -h, --help          Show this help

EXAMPLES:
    $0                              # Update if a newer version is available
    $0 --check                      # Show current vs latest version
    $0 --force                      # Re-install the latest version
    $0 --install-cron               # Set up automatic updates every 25 days
    $0 --install-cron --cron-days 7 # Set up automatic updates every 7 days

REQUIREMENTS:
    - Root privileges
    - code-server installed at SERVDIR (default: /usr/local/code-server)
    - curl

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --check)        CHECK_ONLY=1; shift ;;
            --force)        FORCE=1; shift ;;
            --install-cron) INSTALL_CRON=1; shift ;;
            --cron-days)    CRON_DAYS="$2"; shift 2 ;;
            --installdir)   SERVDIR="$2"; shift 2 ;;
            -h|--help)      show_help ;;
            *) log_err "Unknown option: $1"; show_help ;;
        esac
    done
}

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

preflight_checks() {
    # Root check
    if [[ $EUID -ne 0 ]]; then
        log_err "This script must be run as root"
        exit 1
    fi

    # Required tools
    for cmd in curl tar; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            log_err "Required command not found: $cmd"
            exit 1
        fi
    done

    # code-server must be installed (unless just installing cron)
    if [[ $INSTALL_CRON -eq 0 ]] && ! command -v code-server >/dev/null 2>&1; then
        log_err "code-server is not installed or not in PATH"
        exit 1
    fi
}

# ============================================================================
# GET VERSIONS
# ============================================================================

get_current_version() {
    code-server --version 2>/dev/null | head -1 | awk '{print $1}'
}

get_latest_version() {
    local version
    version=$(curl -fsSL "https://api.github.com/repos/coder/code-server/releases/latest" | \
        grep '"tag_name"' | head -1 | grep -oP 'v\K[0-9]+\.[0-9]+\.[0-9]+')

    if [[ -z "$version" ]]; then
        log_err "Failed to determine latest code-server version"
        exit 1
    fi

    echo "$version"
}

# ============================================================================
# UPDATE CODE-SERVER
# ============================================================================

update_codeserver() {
    local current_version latest_version
    current_version=$(get_current_version)
    latest_version=$(get_latest_version)

    log_info "Installed version: $current_version"
    log_info "Latest version:    $latest_version"

    # Check mode -- just report and exit
    if [[ $CHECK_ONLY -eq 1 ]]; then
        if [[ "$current_version" == "$latest_version" ]]; then
            log "code-server is up to date"
        else
            log_warn "Update available: $current_version -> $latest_version"
        fi
        return
    fi

    # Decide whether to update
    if [[ $FORCE -eq 0 ]]; then
        if [[ "$current_version" == "$latest_version" ]]; then
            log "code-server is already up to date ($current_version)"
            return
        fi

        # Verify the latest version is actually newer (not a downgrade)
        if [[ "$(printf '%s\n' "$current_version" "$latest_version" | sort -V | tail -1)" != "$latest_version" ]]; then
            log "Installed version ($current_version) is newer than latest release ($latest_version)"
            return
        fi
    fi

    log "Updating code-server: $current_version -> $latest_version"

    # Stop service
    systemctl stop code-server

    # Download
    local tarball="code-server-${latest_version}-linux-${ARCH}.tar.gz"
    local url="https://github.com/coder/code-server/releases/download/v${latest_version}/${tarball}"

    cd /tmp
    log_info "Downloading $tarball"
    curl -fsSL -o "$tarball" "$url"
    tar xzf "$tarball"

    # Install
    cp -r "code-server-${latest_version}-linux-${ARCH}"/* "$SERVDIR"/

    # Cleanup
    rm -rf "/tmp/$tarball" "/tmp/code-server-${latest_version}-linux-${ARCH}"

    # Start service
    systemctl start code-server

    log "code-server updated to $latest_version"
}

# ============================================================================
# INSTALL CRON
# ============================================================================

install_cron() {
    local script_dest="/usr/local/bin/update-code-server.sh"
    local script_src

    # Copy this script to /usr/local/bin
    script_src=$(readlink -f "$0")
    if [[ "$script_src" != "$script_dest" ]]; then
        cp "$script_src" "$script_dest"
        log "Copied update script to $script_dest"
    fi
    chmod 700 "$script_dest"

    # Build cron entry -- run every N days at 3:00 AM
    local cron_line="0 3 */${CRON_DAYS} * * ${script_dest} 2>&1 | logger -t update-code-server"
    local cron_marker="# code-server auto-update"

    # Remove existing code-server cron entries
    local existing_cron
    existing_cron=$(crontab -l 2>/dev/null || true)
    local new_cron
    new_cron=$(echo "$existing_cron" | grep -v "$cron_marker" | grep -v "update-code-server" || true)

    # Add new entry
    echo "${new_cron}
${cron_line} ${cron_marker}" | crontab -

    log "Cron job installed -- updates every $CRON_DAYS days at 3:00 AM"
    log_info "View with: crontab -l"
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    parse_args "$@"

    echo ""
    echo "=============================================="
    echo "  Code-Server Update Script v${VERSION}"
    echo "=============================================="
    echo ""

    detect_arch
    preflight_checks

    if [[ $INSTALL_CRON -eq 1 ]]; then
        install_cron
    fi

    # Run update unless we are only installing cron
    if [[ $INSTALL_CRON -eq 0 ]] || command -v code-server >/dev/null 2>&1; then
        update_codeserver
    fi

    echo ""
}

main "$@"