#!/bin/bash
# ============================================================================
# PegaProx Update Script (Archive-based)
# ============================================================================
#
# Downloads the latest release as a tar.gz archive from GitHub.
# This scales to any number of files without needing to list them individually.
#
# Usage:
#   ./update.sh              # Normal update
#   ./update.sh --force      # Force update (skip version check)
#
# Note: Safe to run as root - automatically preserves original file ownership
#
# NS: Rewritten feb 2026 for archive-based updates (code split support)
# ============================================================================

set -e

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

# GitHub URLs
GITHUB_RAW="https://raw.githubusercontent.com/PegaProx/project-pegaprox/main"
# NS: auto-generated by GitHub, no manual release needed
GITHUB_ARCHIVE="https://github.com/PegaProx/project-pegaprox/archive/refs/heads/main.tar.gz"

# Mirror (fallback if GitHub is down)
MIRROR_URL="https://updates.pegaprox.com"
MIRROR_ARCHIVE="https://updates.pegaprox.com/archive/main.tar.gz"

# Find script directory (where PegaProx is installed)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Figure out who owns this installation
ORIGINAL_OWNER=""
if [ -d "config" ]; then
    ORIGINAL_OWNER=$(stat -c '%U:%G' config 2>/dev/null || stat -f '%Su:%Sg' config 2>/dev/null)
elif [ -f "cert.pem" ]; then
    ORIGINAL_OWNER=$(stat -c '%U:%G' cert.pem 2>/dev/null || stat -f '%Su:%Sg' cert.pem 2>/dev/null)
elif [ -d "ssl" ]; then
    ORIGINAL_OWNER=$(stat -c '%U:%G' ssl 2>/dev/null || stat -f '%Su:%Sg' ssl 2>/dev/null)
fi

echo -e "${BLUE}╔════════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║               PegaProx Update Script                       ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════════╝${NC}"
echo ""

# Check if running as root
if [ "$EUID" -eq 0 ]; then
    echo -e "${BLUE}Running as root${NC}"
    if [ -n "$ORIGINAL_OWNER" ]; then
        echo -e "  Will restore ownership to: ${GREEN}$ORIGINAL_OWNER${NC}"
    fi
    echo ""
else
    echo -e "${YELLOW}Tip: sudo ./update.sh for auto service restart${NC}"
    echo ""
fi

# Check current version
CURRENT_VERSION="unknown"
if [ -f "version.json" ]; then
    CURRENT_VERSION=$(grep -o '"version": *"[^"]*"' version.json | cut -d'"' -f4)
fi
echo -e "Current version: ${BLUE}$CURRENT_VERSION${NC}"

# Get latest version (try GitHub first, then mirror)
echo -n "Checking for updates... "
LATEST_VERSION=$(curl -s "$GITHUB_RAW/version.json" 2>/dev/null | grep -o '"version": *"[^"]*"' | cut -d'"' -f4)

if [ -z "$LATEST_VERSION" ]; then
    # GitHub down? try mirror
    LATEST_VERSION=$(curl -s "$MIRROR_URL/version.json" 2>/dev/null | grep -o '"version": *"[^"]*"' | cut -d'"' -f4)
fi

if [ -z "$LATEST_VERSION" ]; then
    echo -e "${RED}Failed${NC}"
    echo "Could not reach GitHub or mirror. Check your internet connection."
    exit 1
fi

echo -e "${GREEN}OK${NC}"
echo -e "Latest version:  ${GREEN}$LATEST_VERSION${NC}"
echo ""

# Compare versions (skip if --force)
if [ "$1" != "--force" ] && [ "$CURRENT_VERSION" == "$LATEST_VERSION" ]; then
    echo -e "${GREEN}✓ You're already on the latest version!${NC}"
    echo ""
    echo "Use ./update.sh --force to re-download anyway"
    exit 0
fi

# Confirm update
echo -e "${YELLOW}Ready to update from $CURRENT_VERSION to $LATEST_VERSION${NC}"
echo ""
read -p "Continue? [y/N] " -n 1 -r
echo ""

if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "Update cancelled."
    exit 0
fi

echo ""
echo -e "${YELLOW}Updating...${NC}"

# Create backup
BACKUP_DIR="backups/backup_${CURRENT_VERSION}_$(date +%Y%m%d_%H%M%S)"
echo -n "Creating backup in $BACKUP_DIR... "
mkdir -p "$BACKUP_DIR"

# Backup important files (not config - that stays)
[ -f "pegaprox_multi_cluster.py" ] && cp pegaprox_multi_cluster.py "$BACKUP_DIR/"
[ -d "pegaprox" ] && cp -r pegaprox "$BACKUP_DIR/"
[ -f "web/index.html" ] && mkdir -p "$BACKUP_DIR/web" && cp web/index.html "$BACKUP_DIR/web/"
[ -f "web/index.html.original" ] && cp web/index.html.original "$BACKUP_DIR/web/"
[ -f "version.json" ] && cp version.json "$BACKUP_DIR/"
[ -f "requirements.txt" ] && cp requirements.txt "$BACKUP_DIR/"

echo -e "${GREEN}OK${NC}"

# Download release archive
echo ""
echo -n "Downloading release archive... "
TMPDIR=$(mktemp -d)
ARCHIVE="$TMPDIR/pegaprox.tar.gz"

if curl -sfL "$GITHUB_ARCHIVE" -o "$ARCHIVE" 2>/dev/null; then
    echo -e "${GREEN}OK (GitHub)${NC}"
elif curl -sfL "$MIRROR_ARCHIVE" -o "$ARCHIVE" 2>/dev/null; then
    echo -e "${GREEN}OK (Mirror)${NC}"
else
    echo -e "${YELLOW}Archive not found, falling back to individual files...${NC}"
    # Fallback: download individual files (for repos without releases)
    # NS: try GitHub first, fall back to mirror
    download_file() {
        local file=$1
        local dir=$(dirname "$file")
        [ "$dir" != "." ] && mkdir -p "$dir"
        echo -n "  $file... "
        if curl -sfL "$GITHUB_RAW/$file" -o "$file.tmp" 2>/dev/null; then
            mv "$file.tmp" "$file"
            echo -e "${GREEN}OK (GitHub)${NC}"
            return 0
        elif curl -sfL "$MIRROR_URL/$file" -o "$file.tmp" 2>/dev/null; then
            mv "$file.tmp" "$file"
            echo -e "${GREEN}OK (mirror)${NC}"
            return 0
        else
            rm -f "$file.tmp"
            echo -e "${RED}FAILED${NC}"
            return 1
        fi
    }

    # Get full file list from version.json (includes pegaprox/ package)
    echo "Fetching file list..."
    PACKAGE_FILES=$(curl -s "$GITHUB_RAW/version.json" 2>/dev/null | python3 -c "
import sys, json
try:
    data = json.load(sys.stdin)
    for f in data.get('update_files', []):
        print(f)
except:
    pass
" 2>/dev/null)

    if [ -z "$PACKAGE_FILES" ]; then
        # Mirror fallback
        PACKAGE_FILES=$(curl -s "$MIRROR_URL/version.json" 2>/dev/null | python3 -c "
import sys, json
try:
    data = json.load(sys.stdin)
    for f in data.get('update_files', []):
        print(f)
except:
    pass
" 2>/dev/null)
    fi

    if [ -n "$PACKAGE_FILES" ]; then
        echo "Downloading ${PACKAGE_FILES##*$'\n'} files..."
        while IFS= read -r pfile; do
            [ -z "$pfile" ] && continue
            download_file "$pfile" || true
        done <<< "$PACKAGE_FILES"
    else
        # absolute fallback - at least get the essentials
        echo "No file list found, downloading essentials..."
        download_file "pegaprox_multi_cluster.py"
        download_file "version.json"
        download_file "requirements.txt"
        download_file "deploy.sh"
        download_file "update.sh"
        download_file "web/index.html"
        download_file "web/index.html.original"
    fi

    rm -rf "$TMPDIR"

    # Skip to pip install
    ARCHIVE=""
fi

# NS Feb 2026 - verify archive integrity via SHA256 checksum
if [ -n "$ARCHIVE" ] && [ -f "$ARCHIVE" ]; then
    echo -n "Verifying archive integrity... "
    SHA_FILE="$TMPDIR/SHA256SUMS"
    SHA_VERIFIED=false
    if curl -sfL "$GITHUB_RAW/SHA256SUMS" -o "$SHA_FILE" 2>/dev/null || \
       curl -sfL "$MIRROR_URL/SHA256SUMS" -o "$SHA_FILE" 2>/dev/null; then
        # SHA256SUMS contains lines like: <hash>  <filename>
        EXPECTED=$(grep -E 'main\.tar\.gz$' "$SHA_FILE" 2>/dev/null | awk '{print $1}')
        if [ -n "$EXPECTED" ]; then
            ACTUAL=$(sha256sum "$ARCHIVE" | awk '{print $1}')
            if [ "$EXPECTED" = "$ACTUAL" ]; then
                echo -e "${GREEN}OK (SHA256 verified)${NC}"
                SHA_VERIFIED=true
            else
                echo -e "${RED}CHECKSUM MISMATCH${NC}"
                echo -e "${RED}Expected: $EXPECTED${NC}"
                echo -e "${RED}Got:      $ACTUAL${NC}"
                echo -e "${RED}Archive may be corrupted or tampered with. Aborting.${NC}"
                rm -rf "$TMPDIR"
                exit 1
            fi
        else
            echo -e "${YELLOW}no matching entry in SHA256SUMS${NC}"
        fi
    else
        echo -e "${YELLOW}SHA256SUMS not available (skipping verification)${NC}"
    fi
fi

# Extract archive if we got one
if [ -n "$ARCHIVE" ] && [ -f "$ARCHIVE" ]; then
    echo -n "Extracting archive... "
    # Extract to temp dir first, then copy (safer)
    EXTRACT_DIR="$TMPDIR/extracted"
    mkdir -p "$EXTRACT_DIR"
    tar xzf "$ARCHIVE" -C "$EXTRACT_DIR" 2>/dev/null

    # Find the actual content (might be in a subdirectory)
    CONTENT_DIR="$EXTRACT_DIR"
    if [ ! -f "$CONTENT_DIR/pegaprox_multi_cluster.py" ]; then
        # Check one level down (GitHub archives often have a subdirectory)
        for subdir in "$EXTRACT_DIR"/*/; do
            if [ -f "${subdir}pegaprox_multi_cluster.py" ]; then
                CONTENT_DIR="$subdir"
                break
            fi
        done
    fi

    if [ -f "$CONTENT_DIR/pegaprox_multi_cluster.py" ]; then
        # Copy files, preserving directory structure
        # Skip: config/, ssl/, logs/, backups/, cert.pem, key.pem, .git/
        if command -v rsync &> /dev/null; then
            rsync -a --exclude='config/' --exclude='ssl/' --exclude='logs/' \
                  --exclude='backups/' --exclude='cert.pem' --exclude='key.pem' \
                  --exclude='.git/' --exclude='.gitignore' \
                  "$CONTENT_DIR/" "$SCRIPT_DIR/"
        else
            # Fallback: cp + tar (works without rsync)
            cd "$CONTENT_DIR"
            tar cf - --exclude='config' --exclude='ssl' --exclude='logs' \
                     --exclude='backups' --exclude='cert.pem' --exclude='key.pem' \
                     --exclude='.git' --exclude='.gitignore' \
                     . | tar xf - -C "$SCRIPT_DIR"
            cd "$SCRIPT_DIR"
        fi
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${RED}FAILED${NC}"
        echo "Archive does not contain pegaprox_multi_cluster.py"
        echo "Restoring from backup..."
        cp "$BACKUP_DIR/pegaprox_multi_cluster.py" . 2>/dev/null || true
        rm -rf "$TMPDIR"
        exit 1
    fi

    rm -rf "$TMPDIR"
fi

# Make scripts executable
chmod +x deploy.sh update.sh 2>/dev/null || true
chmod +x web/Dev/build.sh 2>/dev/null || true

# Fix ownership if running as root
if [ "$EUID" -eq 0 ] && [ -n "$ORIGINAL_OWNER" ] && [ "$ORIGINAL_OWNER" != "root:root" ]; then
    echo -n "Fixing file ownership ($ORIGINAL_OWNER)... "
    chown -R "$ORIGINAL_OWNER" pegaprox_multi_cluster.py version.json requirements.txt 2>/dev/null
    chown -R "$ORIGINAL_OWNER" deploy.sh update.sh 2>/dev/null
    chown -R "$ORIGINAL_OWNER" web/ 2>/dev/null
    [ -d "pegaprox" ] && chown -R "$ORIGINAL_OWNER" pegaprox/ 2>/dev/null
    chown -R "$ORIGINAL_OWNER" backups/ 2>/dev/null
    echo -e "${GREEN}OK${NC}"
fi

# Restore restrictive permissions on config and ssl directories.
# These must be 0700 so that only the service user can read the encrypted
# database and SSL private keys. An update that runs as root via sudo can
# inadvertently leave them world-readable if umask is permissive.
if [ -d "config" ]; then
    chmod 700 config 2>/dev/null || true
fi
if [ -d "config/ssl" ]; then
    chmod 700 config/ssl 2>/dev/null || true
elif [ -d "ssl" ]; then
    chmod 700 ssl 2>/dev/null || true
fi

# Install/update Python packages
echo ""
echo -n "Installing Python packages... "

PIP_SUCCESS=false

if [ -f "venv/bin/python" ] && [ "$PIP_SUCCESS" = false ]; then
    ./venv/bin/python -m pip install -q -r requirements.txt 2>/dev/null && PIP_SUCCESS=true
fi

if [ -f "venv/bin/pip" ] && [ "$PIP_SUCCESS" = false ]; then
    ./venv/bin/pip install -q -r requirements.txt 2>/dev/null && PIP_SUCCESS=true
fi

if [ "$EUID" -eq 0 ] && command -v pip3 &> /dev/null && [ "$PIP_SUCCESS" = false ]; then
    pip3 install -q -r requirements.txt 2>/dev/null && PIP_SUCCESS=true
fi

if command -v pip3 &> /dev/null && [ "$PIP_SUCCESS" = false ]; then
    pip3 install -q --user -r requirements.txt 2>/dev/null && PIP_SUCCESS=true
fi

if command -v python3 &> /dev/null && [ "$PIP_SUCCESS" = false ]; then
    python3 -m pip install -q --user -r requirements.txt 2>/dev/null && PIP_SUCCESS=true
fi

if [ "$PIP_SUCCESS" = true ]; then
    echo -e "${GREEN}OK${NC}"
else
    echo -e "${YELLOW}Couldn't install - run: pip install -r requirements.txt${NC}"
fi

# Restart service
echo ""
echo -n "Restarting PegaProx service... "

if systemctl is-active --quiet pegaprox 2>/dev/null; then
    if systemctl restart pegaprox 2>/dev/null; then
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${YELLOW}Failed - restart manually${NC}"
    fi
elif systemctl is-active --quiet pegaprox.service 2>/dev/null; then
    if systemctl restart pegaprox.service 2>/dev/null; then
        echo -e "${GREEN}OK${NC}"
    else
        echo -e "${YELLOW}Failed - restart manually${NC}"
    fi
else
    echo -e "${YELLOW}No systemd service found${NC}"
    echo "  If running manually, restart with: python3 pegaprox_multi_cluster.py"
fi

# Done!
echo ""
echo -e "${GREEN}╔════════════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║              Update Complete! ✓                            ║${NC}"
echo -e "${GREEN}╚════════════════════════════════════════════════════════════╝${NC}"
echo ""
echo -e "  Updated to version: ${GREEN}$LATEST_VERSION${NC}"
echo -e "  Backup saved to:    ${BLUE}$BACKUP_DIR${NC}"
echo ""
echo "If something went wrong, restore with:"
echo "  cp -r $BACKUP_DIR/* ."
echo ""
