#!/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

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}"

# Backup Configuration
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"

# --- Helper Functions ---

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}"
}

# --- Privilege Check ---

function ensure_root() {
  if [ "$EUID" -ne 0 ]; then
    print_warn "This script requires root privileges."
    print_ok "Attempting to escalate privileges via sudo..."
    
    # Check if sudo is available
    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 replaces the current process, so the script runs anew as root
    exec sudo "$0" "$@"
    exit 0
  fi
  # If we are here, we are root (or sudo-ed successfully)
}

# --- Core Logic ---

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
  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
  fi
  
  print_ok "Running apt update to restore repository state..."
  apt update || true
  
  print_warn "Rollback completed. System restored to previous state."
  print_warn "Backup preserved in: $BACKUP_DIR"
  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..."
  
  mkdir -p "$BACKUP_DIR"
  
  local root_space=$(df / | awk 'NR==2 {print $4}')
  local root_space_mb=$((root_space / 1024))
  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 "Running apt update and upgrade..."
  apt update
  judge "apt update"
  apt upgrade -y
  judge "apt upgrade"
}

function backup_ubuntu_sources() {
  print_ok "Backing up Ubuntu official sources..."
  mkdir -p "$UBUNTU_SOURCE_BACKUP"
  
  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
  
  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"
  
  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
  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"
    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..."
  
  cp /etc/apt/sources.list /etc/apt/sources.list.backup
  
  MIRROR_URL=$(grep -E "^deb " /etc/apt/sources.list | grep -v "security" | head -1 | awk '{print $2}')
  [ -z "$MIRROR_URL" ] && MIRROR_URL="http://archive.ubuntu.com/ubuntu/"
  
  SECURITY_URL=$(grep -E "^deb " /etc/apt/sources.list | grep "security" | head -1 | awk '{print $2}')
  [ -z "$SECURITY_URL" ] && SECURITY_URL="http://security.ubuntu.com/ubuntu/"
  
  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
  
  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..."
  sed -i 's/plucky/questing/g' /etc/apt/sources.list.d/ubuntu.sources
  judge "Replace plucky with questing"
  
  print_ok "Updating package lists..."
  apt update
  judge "apt update with questing"
}

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

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

  print_ok "Running full distribution upgrade..."
  export DEBIAN_FRONTEND=noninteractive
  
  bash -c 'cat > /etc/apt/apt.conf.d/99-local-versions <<EOF
Dpkg::Options {
   "--force-confdef";
   "--force-confold";
}
EOF'
  
  bash -c 'DEBIAN_FRONTEND=noninteractive APT_LISTCHANGES_FRONTEND=none \
  apt-get -y dist-upgrade \
  -o Dpkg::Options::="--force-confdef" \
  -o Dpkg::Options::="--force-confold"'

  judge "apt dist-upgrade"
  rm -f /etc/apt/apt.conf.d/99-local-versions
  unset DEBIAN_FRONTEND
}

function update_release_files() {
  print_ok "Updating release info..."
  
  if [ -f "/etc/os-release" ]; then
    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
  fi
  
  if [ -f "/etc/lsb-release" ]; then
    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
  fi
  judge "Update release files"
}

function restore_ppa_sources() {
  print_ok "Restoring PPA sources..."
  if [ -d "$PPA_BACKUP_DIR" ] && [ "$(ls -A $PPA_BACKUP_DIR)" ]; then
    mv "$PPA_BACKUP_DIR"/* /etc/apt/sources.list.d/
    print_ok "Restored PPAs"
    apt update
    judge "Update after PPA restore"
  else
    print_ok "No PPA sources to restore"
  fi
}

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

  print_ok "Starting AnduinOS upgrade process..."
  echo -e "${Yellow}WARNING: Upgrading from 1.3.7 (plucky) to 1.4.1 (questing).${Font}"
  
  # Since we are likely root now, we check if the script is running interactively before asking
  # If we sudo-ed, we might have lost some TTY context, but usually 'read' works fine.
  if [ -t 0 ]; then
      read -p "Do you want to continue? (y/N): " confirm
      if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
        print_error "Aborted by user."
        exit 1
      fi
  fi
  
  check_disk_space
  update_system
  backup_ubuntu_sources
  backup_and_remove_ppa
  
  if ! detect_apt_format; then
    convert_old_to_new_format
  fi
  
  replace_plucky_with_questing
  install_coreutils_uutils
  run_dist_upgrade
  update_release_files
  restore_ppa_sources
  
  print_ok "Upgrade completed successfully!"
  print_ok "System is now AnduinOS 1.4.1 (questing)"
  print_ok "Backups: $BACKUP_DIR"
  print_warn "Please reboot your system."
}

main