#!/bin/bash

#=================================================
#           AnduinOS Upgrade Script
#=================================================
# This script upgrades AnduinOS from 1.3.7 (plucky)
# to 1.4.1 (questing).
#
# Usage:
# ./do_anduinos_distupgrade.sh
# (Script will auto-elevate to root/sudo)
#=================================================

set -e
set -o pipefail
set -u

# --- 1. Visual Definitions ---
Green="\033[32m"
Red="\033[31m"
Yellow="\033[33m"
Blue="\033[36m"
Font="\033[0m"
OK="${Green}[  OK  ]${Font}"
ERROR="${Red}[FAILED]${Font}"
WARNING="${Yellow}[ WARN ]${Font}"

# --- 2. Configuration ---
# Use /var/backups for persistence and space safety
BACKUP_ROOT="/var/backups/anduinos-upgrade"
BACKUP_DIR="$BACKUP_ROOT/backup_$(date +%Y%m%d_%H%M%S)"
PPA_BACKUP_DIR="$BACKUP_DIR/ppa"
UBUNTU_SOURCE_BACKUP="$BACKUP_DIR/ubuntu_sources"

# --- 3. Helper Functions (Logging) ---

function print_ok() {
  echo -e "${OK} ${Blue} $1 ${Font}"
}

function print_error() {
  echo -e "${ERROR} ${Red} $1 ${Font}"
}

function print_warn() {
  echo -e "${WARNING} ${Yellow} $1 ${Font}"
}

# --- 4. Privilege Check (Auto-Elevation) ---

function ensure_root() {
  if [ "$EUID" -ne 0 ]; then
    print_warn "This script requires root privileges."
    print_ok "Attempting to escalate privileges via sudo..."
    
    if ! command -v sudo &> /dev/null; then
      print_error "sudo is not installed. Please run this script as root."
      exit 1
    fi

    # Re-execute the script with sudo, preserving arguments
    exec sudo "$0" "$@"
    exit 0
  fi
  # If we are here, we are root.
}

# --- 5. Unattended Configuration (Anti-Prompt) ---

function configure_unattended() {
  print_ok "Configuring system for unattended upgrades..."

  # Install debconf-utils to pre-seed answers
  if ! command -v debconf-set-selections &> /dev/null; then
    print_ok "Installing debconf-utils..."
    apt-get update && apt-get install -y debconf-utils
  fi

  # Pre-answer "Yes" to "Restart services during package upgrades without asking?"
  # This kills the libc6/libpam dialog.
  echo '* libraries/restart-without-asking boolean true' | debconf-set-selections

  # Configure environment variables for the session
  # NEEDRESTART_MODE=a : Automatically restart services (fixes purple screen prompt)
  export NEEDRESTART_MODE=a
  export DEBIAN_FRONTEND=noninteractive
  
  judge "Configure unattended mode"
}

# --- 6. Core Logic with Detailed Rollback ---

function rollback_on_error() {
  print_error "An error occurred during the upgrade process"
  print_warn "Starting rollback procedure..."
  
  # Restore ubuntu.sources if backup exists
  if [ -f "$UBUNTU_SOURCE_BACKUP/ubuntu.sources" ]; then
    print_ok "Restoring ubuntu.sources..."
    cp "$UBUNTU_SOURCE_BACKUP/ubuntu.sources" /etc/apt/sources.list.d/
    print_ok "Restored ubuntu.sources"
  fi
  
  # Restore sources.list if backup exists
  if [ -f "$UBUNTU_SOURCE_BACKUP/sources.list" ]; then
    print_ok "Restoring sources.list..."
    cp "$UBUNTU_SOURCE_BACKUP/sources.list" /etc/apt/
    print_ok "Restored sources.list"
  fi
  
  # Restore PPA sources (Detailed Loop from Old Script)
  if [ -d "$PPA_BACKUP_DIR" ]; then
    ppa_count=$(ls -1 "$PPA_BACKUP_DIR" 2>/dev/null | wc -l)
    
    if [ "$ppa_count" -gt 0 ]; then
      print_ok "Restoring PPA sources..."
      for file in "$PPA_BACKUP_DIR"/*; do
        if [ -f "$file" ]; then
          cp "$file" /etc/apt/sources.list.d/
          print_ok "Restored $(basename "$file")"
        fi
      done
    fi
  fi
  
  # Remove temporary apt configuration if exists
  if [ -f "/etc/apt/apt.conf.d/99-local-versions" ]; then
    rm -f /etc/apt/apt.conf.d/99-local-versions
    print_ok "Removed temporary apt configuration"
  fi
  
  # Run apt update to restore repository state
  print_ok "Running apt update to restore repository state..."
  apt update || true
  
  print_warn "Rollback completed"
  print_warn "Your system has been restored to the previous state"
  print_warn "Backup files are preserved in: $BACKUP_DIR"
  print_error "Please check the error messages above and try again"
  
  exit 1
}

function judge() {
  if [[ 0 -eq $? ]]; then
    print_ok "$1 succeeded"
    sleep 0.2
  else
    print_error "$1 failed"
    rollback_on_error
  fi
}

function check_disk_space() {
  print_ok "Checking available disk space..."
  
  # Ensure backup directory exists
  mkdir -p "$BACKUP_DIR"
  
  # Get available space in / (in KB)
  local root_space=$(df / | awk 'NR==2 {print $4}')
  # Convert to MB
  local root_space_mb=$((root_space / 1024))
  # Required space: 2GB
  local required_space=2048
  
  print_ok "Available space in /: ${root_space_mb}MB"
  print_ok "Backup location: $BACKUP_DIR"
  
  if [ "$root_space_mb" -lt "$required_space" ]; then
    print_error "Insufficient disk space in /. Required: ${required_space}MB, Available: ${root_space_mb}MB"
    exit 1
  fi
  
  print_ok "Disk space check passed"
}

function update_system() {
  print_ok "Ensuring current system (1.3 / Ubuntu 25.04) is fully updated..."

  print_ok "Running apt update..."
  apt update
  judge "apt update"

  print_ok "Installing any missing updates for the current version..."
  # Use non-interactive flags to ensure current 1.3 packages are up-to-date
  DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=a apt upgrade -y
  judge "apt upgrade"
}

function backup_ubuntu_sources() {
  print_ok "Backing up Ubuntu official sources..."
  
  mkdir -p "$UBUNTU_SOURCE_BACKUP"
  
  # Backup ubuntu.sources if exists
  if [ -f "/etc/apt/sources.list.d/ubuntu.sources" ]; then
    cp /etc/apt/sources.list.d/ubuntu.sources "$UBUNTU_SOURCE_BACKUP/"
    print_ok "Backed up ubuntu.sources"
  fi
  
  # Backup sources.list if it exists and is not empty
  if [ -f "/etc/apt/sources.list" ] && [ -s "/etc/apt/sources.list" ]; then
    cp /etc/apt/sources.list "$UBUNTU_SOURCE_BACKUP/"
    print_ok "Backed up sources.list"
  fi
  
  judge "Backup Ubuntu sources"
}

function backup_and_remove_ppa() {
  print_ok "Backing up and temporarily removing PPA sources..."
  
  mkdir -p "$PPA_BACKUP_DIR"
  
  # Move all files in /etc/apt/sources.list.d/ except ubuntu.sources
  if [ -d "/etc/apt/sources.list.d" ]; then
    for file in /etc/apt/sources.list.d/*; do
      if [ -f "$file" ] && [ "$(basename "$file")" != "ubuntu.sources" ]; then
        mv "$file" "$PPA_BACKUP_DIR/"
        print_ok "Moved $(basename "$file") to backup"
      fi
    done
  fi
  
  print_ok "PPA sources moved to: $PPA_BACKUP_DIR"
  judge "Backup and remove PPA sources"
}

function detect_apt_format() {
  print_ok "Detecting APT source format..."
  
  if [ -f "/etc/apt/sources.list.d/ubuntu.sources" ]; then
    print_ok "Detected new DEB822 format (ubuntu.sources)"
    return 0
  elif [ -f "/etc/apt/sources.list" ] && [ -s "/etc/apt/sources.list" ]; then
    print_ok "Detected old format (sources.list)"
    return 1
  else
    print_error "Cannot detect APT source format"
    rollback_on_error
  fi
}

function convert_old_to_new_format() {
  print_ok "Converting old format to new DEB822 format..."
  
  if [ ! -f "/etc/apt/sources.list" ] || [ ! -s "/etc/apt/sources.list" ]; then
    print_error "/etc/apt/sources.list not found or empty"
    rollback_on_error
  fi
  
  # Backup the old sources.list
  cp /etc/apt/sources.list /etc/apt/sources.list.backup
  
  # Extract mirror URLs from existing sources.list (Restored from Old Script)
  # Try to detect main archive mirror
  MIRROR_URL=$(grep -E "^deb " /etc/apt/sources.list | grep -v "security" | head -1 | awk '{print $2}')
  if [ -z "$MIRROR_URL" ]; then
    MIRROR_URL="http://archive.ubuntu.com/ubuntu/"
  fi
  
  # Try to detect security mirror
  SECURITY_URL=$(grep -E "^deb " /etc/apt/sources.list | grep "security" | head -1 | awk '{print $2}')
  if [ -z "$SECURITY_URL" ]; then
    SECURITY_URL="http://security.ubuntu.com/ubuntu/"
  fi
  
  print_ok "Detected mirror URL: $MIRROR_URL"
  print_ok "Detected security URL: $SECURITY_URL"
  
  # Create new ubuntu.sources file using detected mirrors
  bash -c "cat > /etc/apt/sources.list.d/ubuntu.sources" <<EOF
Types: deb
URIs: ${MIRROR_URL}
Suites: plucky plucky-updates plucky-backports
Components: main restricted universe multiverse
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg

Types: deb
URIs: ${SECURITY_URL}
Suites: plucky-security
Components: main restricted universe multiverse
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg
EOF
  
  # Clear the old sources.list
  bash -c 'echo "# This file is deprecated. See /etc/apt/sources.list.d/ubuntu.sources" > /etc/apt/sources.list'
  
  judge "Convert to new format"
}

function replace_plucky_with_questing() {
  print_ok "Replacing plucky with questing in ubuntu.sources..."
  
  if [ ! -f "/etc/apt/sources.list.d/ubuntu.sources" ]; then
    print_error "/etc/apt/sources.list.d/ubuntu.sources not found"
    rollback_on_error
  fi
  
  sed -i 's/plucky/questing/g' /etc/apt/sources.list.d/ubuntu.sources
  judge "Replace plucky with questing"
  
  print_ok "Running apt update with questing repositories..."
  apt update
  judge "apt update with questing"
}

function install_coreutils_uutils() {
  print_ok "Installing coreutils-from-uutils..."
  
  DEBIAN_FRONTEND=noninteractive apt install -y coreutils-from-uutils
  judge "Install coreutils-from-uutils"
}

function run_dist_upgrade() {
  print_ok "Simulating apt dist-upgrade first..."
  apt -s dist-upgrade > /dev/null
  judge "apt -s dist-upgrade"

  print_ok "Running apt dist-upgrade in non-interactive mode..."
  
  # Configure dpkg to keep local versions by default
  bash -c 'cat > /etc/apt/apt.conf.d/99-local-versions <<EOF
Dpkg::Options {
   "--force-confdef";
   "--force-confold";
}
EOF'
  
  # Run dist-upgrade
  # Combine variables: DEBIAN_FRONTEND, NEEDRESTART_MODE, APT_LISTCHANGES
  bash -c 'DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=a APT_LISTCHANGES_FRONTEND=none \
  apt-get -y dist-upgrade \
  -o Dpkg::Options::="--force-confdef" \
  -o Dpkg::Options::="--force-confold"'

  judge "apt dist-upgrade"
  
  # Remove temporary configuration
  rm -f /etc/apt/apt.conf.d/99-local-versions
}

function update_release_files() {
  print_ok "Updating release information files to 1.4.1..."
  
  # Update /etc/os-release
  if [ -f "/etc/os-release" ]; then
    print_ok "Updating /etc/os-release..."
    bash -c "cat > /etc/os-release" <<EOF
PRETTY_NAME="AnduinOS 1.4.1"
NAME="AnduinOS"
VERSION_ID="1.4.1"
VERSION="1.4.1 (questing)"
VERSION_CODENAME=questing
ID=ubuntu
ID_LIKE=debian
HOME_URL="https://www.anduinos.com/"
SUPPORT_URL="https://github.com/Anduin2017/AnduinOS/discussions"
BUG_REPORT_URL="https://github.com/Anduin2017/AnduinOS/issues"
PRIVACY_POLICY_URL="https://www.ubuntu.com/legal/terms-and-policies/privacy-policy"
UBUNTU_CODENAME=questing
EOF

    judge "Update /etc/os-release"
  fi
  
  # Update /etc/lsb-release
  if [ -f "/etc/lsb-release" ]; then
    print_ok "Updating /etc/lsb-release..."
    
    bash -c "cat > /etc/lsb-release" <<EOF
DISTRIB_ID=AnduinOS
DISTRIB_RELEASE=1.4.1
DISTRIB_CODENAME=questing
DISTRIB_DESCRIPTION="AnduinOS 1.4.1"
EOF

    judge "Update /etc/lsb-release"
  fi
  
  print_ok "Release files updated successfully"
}

function restore_ppa_sources() {
  print_ok "Restoring PPA sources..."
  
  if [ -d "$PPA_BACKUP_DIR" ]; then
    ppa_count=$(ls -1 "$PPA_BACKUP_DIR" 2>/dev/null | wc -l)
    
    if [ "$ppa_count" -gt 0 ]; then
      for file in "$PPA_BACKUP_DIR"/*; do
        if [ -f "$file" ]; then
          mv "$file" /etc/apt/sources.list.d/
          print_ok "Restored $(basename "$file")"
        fi
      done
      
      print_ok "Running apt update with restored PPAs..."
      apt update
      judge "Restore PPA sources and update"
    else
      print_ok "No PPA sources to restore"
    fi
  else
    print_warn "PPA backup directory not found, skipping restore"
  fi
}

function main() {
  # 1. Ensure we are root first
  ensure_root

  print_ok "Starting AnduinOS upgrade process..."
  
  echo -e "${Yellow}WARNING: This script will upgrade your system from 1.3.7 (plucky) to 1.4.1 (questing).${Font}"
  echo -e "${Yellow}Please ensure you have backed up important data before proceeding.${Font}"
  
  # Interactive check only if we have a terminal (TTY)
  if [ -t 0 ]; then
      read -p "Do you want to continue? (y/N): " confirm
      if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
        print_error "Upgrade process aborted by user."
        exit 1
      fi
  fi
  
  # Step 0: Configure Unattended (Anti-Prompt)
  configure_unattended

  # Step 1: Check disk space
  check_disk_space
  
  # Step 2: Update current system
  update_system
  
  # Step 3: Backup Ubuntu official sources
  backup_ubuntu_sources
  
  # Step 4: Backup and remove PPA sources
  backup_and_remove_ppa
  
  # Step 5: Detect and convert APT format if needed
  if ! detect_apt_format; then
    convert_old_to_new_format
  fi
  
  # Step 6: Replace plucky with questing
  replace_plucky_with_questing
  
  # Step 7: Install coreutils-from-uutils
  install_coreutils_uutils
  
  # Step 8: Run dist-upgrade
  run_dist_upgrade
  
  # Step 9: Update release files (to 1.4.1)
  update_release_files
  
  # Step 10: Restore PPA sources
  restore_ppa_sources
  
  print_ok "Upgrade completed successfully!"
  print_ok "Your system has been upgraded to AnduinOS 1.4.1 (questing)"
  print_ok "Backup files are stored in: $BACKUP_DIR"
  print_warn "Please reboot your system to complete the upgrade."
}

main