#!/bin/bash
#==========================
# Set up the environment
#==========================
set -e                  # exit on error
set -o pipefail         # exit on pipeline error
set -u                  # treat unset variable as error

#==========================
# Color
#==========================
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}"

#==========================
# Print Colorful Text
#==========================
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}"
}

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

# Clean up function
#==========================
function clean_up() {
  # Clean up
  print_ok "Cleaning up old files..."
  sudo umount /mnt/anduinos_squashfs >/dev/null 2>&1 || true
  sudo umount /mnt/anduinos_iso >/dev/null 2>&1 || true
  sudo rm -rf /mnt/anduinos_squashfs >/dev/null 2>&1 || true
  sudo rm -rf /mnt/anduinos_iso >/dev/null 2>&1 || true
  sudo rm /tmp/AnduinOS-1.4.0* >/dev/null 2>&1 || true
  judge "Cleanup"
}

clean_up

print_ok "Checking system compatibility..."
codename=$(lsb_release -cs)
if [[ "$codename" != "questing" ]] then
    print_error "This upgrade script can only be run *from* AnduinOS Questing."
    exit 1
fi
judge "System compatibility check"

print_ok "Installing required packages (aria2, curl, lsb-release)..."
sudo apt install -y aria2 curl lsb-release
judge "Install required packages"

# $LANG maybe en_US.UTF-8, so we use only the `en_US` part
CURRENT_LANG=${LANG%%.*}
DOWNLOAD_URL="https://download.anduinos.com/1.4/1.4.0/AnduinOS-1.4.0-${CURRENT_LANG}.torrent"
HASH_URL="https://download.anduinos.com/1.4/1.4.0/AnduinOS-1.4.0-${CURRENT_LANG}.sha256"

print_ok "Current system language detected: ${CURRENT_LANG}"
print_ok "Attempting to download with URL: ${DOWNLOAD_URL}"

# If download URL is not reachable, fallback to en_US
if ! curl --head --silent --fail "$DOWNLOAD_URL" >/dev/null; then
    print_warn "Language pack for ${CURRENT_LANG} not found, falling back to en_US"
    DOWNLOAD_URL="https://download.anduinos.com/1.4/1.4.0/AnduinOS-1.4.0-en_US.torrent"
    HASH_URL="https://download.anduinos.com/1.4/1.4.0/AnduinOS-1.4.0-en_US.sha256"
fi

# Ensure download URL is reachable
if ! curl --head --silent --fail "$DOWNLOAD_URL" >/dev/null; then
    print_error "Download URL is not reachable. Please check your network connection."
    exit 1
fi

print_ok "Downloading AnduinOS 1.4.0 torrent, please wait..."
curl -o /tmp/AnduinOS-1.4.0.torrent "$DOWNLOAD_URL"
curl -o /tmp/AnduinOS-1.4.0.sha256 "$HASH_URL"
judge "Download AnduinOS 1.4.0 torrent"

print_ok "Starting download via aria2..."
aria2c --allow-overwrite=true --seed-ratio=0.0 --seed-time=0 -x 16 -s 16 -k 1M -d /tmp /tmp/AnduinOS-1.4.0.torrent
judge "Download AnduinOS 1.4.0 ISO"

ISO_FILE_PATH=$(ls /tmp/AnduinOS-1.4.0*.iso | head -n 1)
print_ok "Ensure downloaded ISO file exists..."
if [[ -f "$ISO_FILE_PATH" ]]; then
    print_ok "Downloaded ISO file found: $ISO_FILE_PATH"
else
    print_error "Downloaded ISO file not found."
    exit 1
fi

SHA256_FILE_PATH="/tmp/AnduinOS-1.4.0.sha256"

print_ok "Verifying download integrity..."
ACTUAL_SHA256=$(sha256sum "$ISO_FILE_PATH" | awk '{print $1}')
EXPECTED_SHA256=$(grep 'SHA256:' "$SHA256_FILE_PATH" | awk '{print $2}')
if [[ "$ACTUAL_SHA256" == "$EXPECTED_SHA256" ]]; then
    print_ok "SHA256 checksum verification passed."
else
    print_ok "Expected SHA256: $EXPECTED_SHA256"
    print_ok "Actual SHA256:   $ACTUAL_SHA256"
    print_error "SHA256 checksum verification failed. The downloaded file may be corrupted."
    exit 1
fi

print_ok "Mounting the ISO..."
sudo mkdir -p /mnt/anduinos_iso
sudo mount -o loop,ro "$ISO_FILE_PATH" /mnt/anduinos_iso
judge "Mount ISO"

print_ok "Verifying content in the ISO..."
(cd /mnt/anduinos_iso && sudo md5sum -c md5sum.txt)
judge "ISO content integrity verification"

print_ok "Mounting the filesystem.squashfs..."
sudo mkdir -p /mnt/anduinos_squashfs
sudo mount -o loop,ro /mnt/anduinos_iso/casper/filesystem.squashfs /mnt/anduinos_squashfs
judge "Mount filesystem.squashfs"

print_ok "Updating package mirrors..."
curl -s https://gitlab.aiursoft.cn/anduin/init-server/-/raw/master/mirror.sh?ref_type=heads | bash
sudo apt update
judge "Update package mirrors"

# Reset apt settings
print_ok "Resetting APT configuration files..."
sudo rm /etc/apt/preferences.d/* >/dev/null 2>&1 || true
judge "Reset APT configuration files"

# Add Mozilla Team PPA
print_ok "Updating Mozilla Team PPA..."
sudo rm -f /etc/apt/sources.list.d/mozillateam*
sudo rsync -Aax /mnt/anduinos_squashfs/etc/apt/sources.list.d/mozillateam* /etc/apt/sources.list.d/
judge "Update Mozilla Team PPA"

# Install missing package.
print_ok "Installing missing packages from the new release..."
MANIFEST_FILE="/mnt/anduinos_iso/casper/filesystem.manifest-desktop"
cut -d' ' -f1 "$MANIFEST_FILE" \
  | grep -v '^linux-' \
  | grep -v '^lib' \
  | grep -v '^plymouth-' \
  | grep -v '^software-properties-' \
  | xargs sudo apt install -y  # <-- 让 xargs 自己调用 apt
judge "Install missing packages"

# Remove obsolete packages.
print_ok "Removing obsolete packages..."
sudo apt autoremove -y \
  distro-info \
  software-properties-gtk \
  ubuntu-advantage-tools \
  ubuntu-pro-client \
  ubuntu-pro-client-l10n \
  ubuntu-release-upgrader-gtk \
  ubuntu-report \
  ubuntu-settings \
  update-notifier-common \
  update-manager \
  update-manager-core \
  update-notifier \
  ubuntu-release-upgrader-core \
  ubuntu-advantage-desktop-daemon \
  kgx
judge "Remove obsolete packages"

print_ok "Upgrading installed packages..."
sudo apt upgrade -y
sudo apt autoremove --purge -y
judge "System package cleanup"

print_ok "Upgrading GNOME Shell extensions..."
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/usr/share/gnome-shell/extensions/ /usr/share/gnome-shell/extensions/
judge "Upgrade GNOME Shell extensions"

print_ok "Upgrading icon and theme files..."
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/usr/share/icons/ /usr/share/icons/
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/usr/share/themes/ /usr/share/themes/
judge "Upgrade icon and theme files"

print_ok "Upgrading desktop backgrounds..."
sudo rsync -Aax --update /mnt/anduinos_squashfs/usr/share/backgrounds/ /usr/share/backgrounds/
sudo rsync -Aax --update /mnt/anduinos_squashfs/usr/share/gnome-background-properties/ /usr/share/gnome-background-properties/
judge "Upgrade desktop backgrounds"

print_ok "Upgrading APT configuration files..."
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/etc/apt/apt.conf.d/ /etc/apt/apt.conf.d/
judge "Upgrade APT configuration files"

print_ok "Upgrading APT preferences files..."
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/etc/apt/preferences.d/ /etc/apt/preferences.d/
judge "Upgrade APT preferences files"

print_ok "Upgrading session files..."
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/usr/share/gnome-session/sessions/ /usr/share/gnome-session/sessions/
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/usr/share/wayland-sessions/ /usr/share/wayland-sessions/
judge "Upgrade session files"

print_ok "Upgrading pixmaps..."
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/usr/share/pixmaps/ /usr/share/pixmaps/
judge "Upgrade pixmaps"

print_ok "Upgrading python-apt templates and distro info..."
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/usr/share/python-apt/templates/ /usr/share/python-apt/templates/
sudo rsync -Aax --update --delete /mnt/anduinos_squashfs/usr/share/distro-info/ /usr/share/distro-info/
judge "Upgrade python-apt templates and distro info"

print_ok "Upgrading deskmon service..."
sudo rsync -Aax /mnt/anduinos_squashfs/usr/local/bin/deskmon /usr/local/bin/deskmon
sudo rsync -Aax /mnt/anduinos_squashfs/etc/systemd/user/deskmon.service /etc/systemd/user/deskmon.service
sudo rsync -Aax /mnt/anduinos_squashfs/etc/systemd/user/default.target.wants/deskmon.service /etc/systemd/user/default.target.wants/deskmon.service
judge "Upgrade deskmon service"

print_ok "Updating system version information..."
sudo rsync -Aax /mnt/anduinos_squashfs/usr/local/bin/do_anduinos_upgrade /usr/local/bin/do_anduinos_upgrade
sudo rsync -Aax /mnt/anduinos_squashfs/usr/bin/add-apt-repository /usr/bin/add-apt-repository
sudo rsync -Aax /mnt/anduinos_squashfs/etc/lsb-release /etc/lsb-release
sudo rsync -Aax /mnt/anduinos_squashfs/etc/issue /etc/issue
sudo rsync -Aax /mnt/anduinos_squashfs/etc/issue.net /etc/issue.net
sudo rsync -Aax /mnt/anduinos_squashfs/etc/os-release /etc/os-release
sudo rsync -Aax /mnt/anduinos_squashfs/usr/lib/os-release /usr/lib/os-release
sudo rsync -Aax /mnt/anduinos_squashfs/etc/legal /etc/legal
sudo rsync -Aax /mnt/anduinos_squashfs/etc/sysctl.d/20-apparmor-donotrestrict.conf /etc/sysctl.d/20-apparmor-donotrestrict.conf
judge "Update system version information"

print_ok "Applying dconf settings patch..."
VERSION=$(grep -oP "VERSION_ID=\"\\K\\d+\\.\\d+" /etc/os-release)
PATCH_URL="https://gitlab.aiursoft.cn/anduin/anduinos/-/raw/$VERSION/src/mods/35-dconf-patch/dconf.ini?ref_type=heads"
COMMAND_TO_RUN="set -o pipefail; curl -sL '$PATCH_URL' | dconf load /org/gnome/"
if [[ -n "${SUDO_USER-}" ]]; then
    print_ok "Applying dconf patch for user: $SUDO_USER"
    sudo -u "$SUDO_USER" DBUS_SESSION_BUS_ADDRESS="unix:path=/run/user/$(id -u $SUDO_USER)/bus" \
    bash -c "$COMMAND_TO_RUN"
else
    print_warn "Running as root without sudo context. Skipping user dconf patch."
    true
fi
judge "Apply dconf settings patch"

print_ok "Upgrade completed! Please reboot your system to apply all changes."

print_ok "Starting cleanup..."
clean_up