#!/bin/bash

#=================================================
#           AnduinOS Upgrade Script
#=================================================
# This script upgrades AnduinOS 
# to 1.4.0 (questing).
#
# Run this script as root.
# or user with privileges.
#
# Example:
# sudo ./do_anduinos_distupgrade.sh
#=================================================

set -e
set -o pipefail
set -u

Green="\033[32m"
Red="\033[31m"
Yellow="\033[33m"
Blue="\033[36m"
Font="\033[0m"
GreenBG="\033[42;37m"
RedBG="\033[41;37m"
OK="${Green}[  OK  ]${Font}"
ERROR="${Red}[FAILED]${Font}"
WARNING="${Yellow}[ WARN ]${Font}"

BACKUP_DIR="/tmp/anduinos_upgrade_backup_$(date +%Y%m%d_%H%M%S)"
PPA_BACKUP_DIR="$BACKUP_DIR/ppa"
UBUNTU_SOURCE_BACKUP="$BACKUP_DIR/ubuntu_sources"

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

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
    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..."
  
  # Get available space in /tmp (in KB)
  local tmp_space=$(df /tmp | awk 'NR==2 {print $4}')
  # Get available space in / (in KB)
  local root_space=$(df / | awk 'NR==2 {print $4}')
  
  # Convert to MB
  local tmp_space_mb=$((tmp_space / 1024))
  local root_space_mb=$((root_space / 1024))
  
  # Required space: 2GB = 2048MB
  local required_space=2048
  
  print_ok "Available space in /tmp: ${tmp_space_mb}MB"
  print_ok "Available space in /: ${root_space_mb}MB"
  
  if [ "$tmp_space_mb" -lt "$required_space" ]; then
    print_error "Insufficient disk space in /tmp. Required: ${required_space}MB, Available: ${tmp_space_mb}MB"
    exit 1
  fi
  
  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"
  
  # 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
  # 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: questing questing-updates questing-backports
Components: main restricted universe multiverse
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg

Types: deb
URIs: ${SECURITY_URL}
Suites: questing-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
  
  # todo: replace 1.2.x(oracular)  
  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..."
  
  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
  judge "apt -s dist-upgrade"

  print_ok "Running apt dist-upgrade in non-interactive mode..."
  
  # Set environment for non-interactive mode
  export DEBIAN_FRONTEND=noninteractive
  
  # 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
  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"
  
  # Remove temporary configuration
  rm -f /etc/apt/apt.conf.d/99-local-versions
  
  unset DEBIAN_FRONTEND
}

function update_release_files() {
  print_ok "Updating release information files..."
  
  # 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.0"
NAME="AnduinOS"
VERSION_ID="1.4.0"
VERSION="1.4.0 (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.0
DISTRIB_CODENAME=questing
DISTRIB_DESCRIPTION="AnduinOS 1.4.0"
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() {
  print_ok "Starting AnduinOS upgrade process..."
  
  echo -e "${Yellow}WARNING: This script will upgrade your system to 1.4.0 (questing).${Font}"
  echo -e "${Yellow}Please ensure you have backed up important data before proceeding.${Font}"
  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
  
  # Check if running as root
  if [[ "$(id -u)" -ne 0 ]]; then
    print_error "This script must be run as root."
    exit 1
  fi
  
  # Step 0: Check disk space
  check_disk_space
  
  # Step 1: Update current system
  update_system
  
  # Step 2: Backup Ubuntu official sources
  backup_ubuntu_sources
  
  # Step 3: Backup and remove PPA sources
  backup_and_remove_ppa
  
  # Step 4: Detect and convert APT format if needed
  if ! detect_apt_format; then
    convert_old_to_new_format
  fi
  
  # Step 5: Replace plucky with questing
  replace_plucky_with_questing
  
  # Step 6: Install coreutils-from-uutils
  install_coreutils_uutils
  
  # Step 7: Run dist-upgrade
  run_dist_upgrade
  
  # Step 8: Update release files
  update_release_files
  
  # Step 9: Restore PPA sources
  restore_ppa_sources
  
  print_ok "Upgrade completed successfully!"
  print_ok "Your system has been upgraded to AnduinOS 1.4.0 (questing)"
  print_ok "Backup files are stored in: $BACKUP_DIR"
  print_warn "Please reboot your system to complete the upgrade."
}

main