#!/bin/bash

# ==============================================================================
#  SonicDE Artix Linux Installer (Pure Artix Edition)
#  Features: No Arch Repos, Native Pacman Install, OpenRC, Multi-Bootloader
# ==============================================================================

# -- 0. Root Check --
if [ "$EUID" -ne 0 ]; then
  echo "Error: This script must be run as root. Try: sudo ./install.sh"
  exit 1
fi

# -- Colors & Variables --
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m'
MOUNT_POINT="/mnt"
LOG="/tmp/installer.log"

[ -f "$LOG" ] && rm "$LOG"

# -- Helper Functions --
log() { echo -e "${GREEN}[*] $1${NC}" | tee -a "$LOG"; }
err() { echo -e "${RED}[!] $1${NC}" | tee -a "$LOG"; exit 1; }

# -- 1. Pre-Flight Checks --
pre_flight() {
    clear
    log "Checking environment..."
    if ! ping -c 1 google.com &>/dev/null; then
        err "No internet connection. Please connect manually."
    fi

    log "Updating Artix Keyring..."
    pacman -Sy --noconfirm artix-keyring
    pacman -S --noconfirm dialog git || err "Failed to install dependencies."
}

# -- 2. Configure Artix Repos (Enable Galaxy) --
setup_artix_repos() {
    local TARGET_ROOT="$1"
    local CONF_PATH="${TARGET_ROOT}/etc/pacman.conf"

    log "Configuring Artix Repositories on ${TARGET_ROOT:-Live System}..."

    # Ensure [galaxy] is enabled (often needed for community packages)
    if ! grep -q "^\[galaxy\]" "$CONF_PATH"; then
        # Try to uncomment if it exists
        sed -i '/#\[galaxy\]/s/^#//g' "$CONF_PATH" 2>/dev/null
        sed -i '/\[galaxy\]/{n;s/^#//}' "$CONF_PATH" 2>/dev/null
        
        # If still missing, append it
        if ! grep -q "^\[galaxy\]" "$CONF_PATH"; then
            echo -e "\n[galaxy]\nInclude = /etc/pacman.d/mirrorlist" >> "$CONF_PATH"
        fi
    fi

    # Ensure [lib32] is enabled (often needed for Steam/Wine/Gaming)
    if ! grep -q "^\[lib32\]" "$CONF_PATH"; then
        sed -i '/#\[lib32\]/s/^#//g' "$CONF_PATH" 2>/dev/null
        sed -i '/\[lib32\]/{n;s/^#//}' "$CONF_PATH" 2>/dev/null
    fi

    # Refresh DB
    if [ -z "$TARGET_ROOT" ]; then pacman -Sy; fi
}

# -- 3. Menus --
select_menu_options() {
    # Drive
    OPTIONS=()
    while read -r name size model; do
        desc="${size}_${model// /_}"
        OPTIONS+=("$name" "$desc")
    done < <(lsblk -dnp -o NAME,SIZE,MODEL | grep -v 'loop' | grep -v 'sr0')

    if [ ${#OPTIONS[@]} -eq 0 ]; then err "No disks found."; fi
    DISK=$(dialog --stdout --menu "Select Installation Disk" 15 60 5 "${OPTIONS[@]}") || exit 1
    
    # Partition Method
    # Note: Removed tparted because it is an AUR package. 
    # Without Arch support, we default to standard tools (cfdisk).
    METHOD=$(dialog --stdout --menu "Partitioning Method" 12 60 2 \
        "Auto" "Automatic Btrfs (Wipe Disk)" \
        "Manual" "Manual via cfdisk (Standard)") || exit 1

    # Bootloader
    BOOTLOADER=$(dialog --stdout --menu "Select Bootloader" 13 60 3 \
        "GRUB" "Standard, reliable" \
        "rEFInd" "Graphical, auto-detects" \
        "Limine" "Modern, simple config") || exit 1

    # Timezone
    REGIONS=$(find /usr/share/zoneinfo -maxdepth 1 -mindepth 1 -type d -printf "%f\n" | sort)
    REGION_ARGS=(); for r in $REGIONS; do REGION_ARGS+=("$r" "-"); done
    REGION=$(dialog --stdout --menu "Select Region" 15 40 5 "${REGION_ARGS[@]}") || exit 1
    
    CITIES=$(find "/usr/share/zoneinfo/$REGION" -maxdepth 1 -mindepth 1 -type f -printf "%f\n" | sort)
    CITY_ARGS=(); for c in $CITIES; do CITY_ARGS+=("$c" "-"); done
    CITY=$(dialog --stdout --menu "Select City" 15 40 10 "${CITY_ARGS[@]}") || exit 1
    TIMEZONE="${REGION}/${CITY}"
}

# -- 4. Partitioning --
partition_drive() {
    if [ "$METHOD" == "Auto" ]; then
        dialog --defaultno --yesno "WARNING: ALL DATA ON $DISK WILL BE ERASED. PROCEED?" 10 60 || exit 1
        clear
        log "Wiping $DISK..."
        sgdisk -Z "$DISK"
        sgdisk -n 1:0:+512M -t 1:ef00 -c 1:"EFI" "$DISK"
        sgdisk -n 2:0:+4G   -t 2:8200 -c 2:"SWAP" "$DISK"
        sgdisk -n 3:0:0     -t 3:8300 -c 3:"ROOT" "$DISK"
        
        if [[ "$DISK" == *"nvme"* ]]; then P1="${DISK}p1"; P2="${DISK}p2"; P3="${DISK}p3"; else P1="${DISK}1"; P2="${DISK}2"; P3="${DISK}3"; fi

        mkfs.vfat -F32 "$P1"
        mkswap "$P2" && swapon "$P2"
        mkfs.btrfs -f "$P3"

        mount "$P3" "$MOUNT_POINT"
        btrfs subvolume create "$MOUNT_POINT/@"
        btrfs subvolume create "$MOUNT_POINT/@home"
        btrfs subvolume create "$MOUNT_POINT/@snapshots"
        btrfs subvolume create "$MOUNT_POINT/@var_log"
        umount "$MOUNT_POINT"

        mount -o noatime,compress=zstd,subvol=@ "$P3" "$MOUNT_POINT"
        mkdir -p "$MOUNT_POINT"/{home,.snapshots,var/log,boot}
        mount -o noatime,compress=zstd,subvol=@home "$P3" "$MOUNT_POINT/home"
        mount -o noatime,compress=zstd,subvol=@snapshots "$P3" "$MOUNT_POINT/.snapshots"
        mount -o noatime,compress=zstd,subvol=@var_log "$P3" "$MOUNT_POINT/var/log"
        mount "$P1" "$MOUNT_POINT/boot"
    else
        clear
        log "Launching cfdisk..."
        cfdisk "$DISK"
        
        dialog --msgbox "Partitioning complete.\n\nPlease manually mount partitions to /mnt now.\n(Alt+F2 to open terminal, then 'exit' to return)" 10 60
        dialog --yesno "Are partitions mounted at /mnt?" 10 60 || exit 1
    fi
}

# -- 5. Base Install --
install_base() {
    clear
    log "Installing Artix Base System..."
    # REMOVED: artix-archlinux-support, yaourtix, pikaur
    # ADDED: sonicde-meta (Assuming it is in Artix repos as requested)
    # We install base system first, SonicDE is installed in config step to catch errors cleanly.
    
    basestrap "$MOUNT_POINT" base base-devel openrc elogind elogind-openrc \
                             linux linux-firmware networkmanager networkmanager-openrc \
                             nano git btrfs-progs || err "Basestrap failed."
                             
    fstabgen -U "$MOUNT_POINT" >> "$MOUNT_POINT/etc/fstab"

    if [ ! -f "$MOUNT_POINT/bin/bash" ]; then
        err "Base install failed. /mnt/bin/bash not found."
    fi
}

# -- 6. Configure Target (Chroot) --
configure_system() {
    HOSTNAME=$(dialog --stdout --inputbox "Enter Hostname:" 10 40 "artixlinux")
    USERNAME=$(dialog --stdout --inputbox "Enter Username:" 10 40 "user")
    PASSWORD=$(dialog --stdout --passwordbox "Enter Password:" 10 40)

    # Prepare Chroot
    cp /etc/resolv.conf "$MOUNT_POINT/etc/resolv.conf"
    
    # Generate Script
    cat <<EOF > "$MOUNT_POINT/setup_internal.sh"
#!/bin/bash
set -e
exec < /dev/tty

echo "------------------------------------------------"
echo "Starting Artix Configuration..."
echo "------------------------------------------------"

# 1. System Config
ln -sf /usr/share/zoneinfo/$TIMEZONE /etc/localtime
hwclock --systohc
echo "en_US.UTF-8 UTF-8" > /etc/locale.gen
locale-gen
echo "LANG=en_US.UTF-8" > /etc/locale.conf
echo "export LANG=en_US.UTF-8" >> /etc/profile
echo "$HOSTNAME" > /etc/hostname
echo "127.0.0.1 localhost" >> /etc/hosts
echo "127.0.1.1 $HOSTNAME.localdomain $HOSTNAME" >> /etc/hosts

# 2. OpenRC Networking
rc-update add NetworkManager default

# 3. Setup Repos (Target)
pacman-key --init
pacman-key --populate artix

# Enable [galaxy] in target if missing
if ! grep -q "^\[galaxy\]" /etc/pacman.conf; then
    echo -e "\n[galaxy]\nInclude = /etc/pacman.d/mirrorlist" >> /etc/pacman.conf
fi
pacman -Sy

# 4. User Setup
echo "root:$PASSWORD" | chpasswd
useradd -m -G wheel -s /bin/bash "$USERNAME" || true
echo "$USERNAME:$PASSWORD" | chpasswd
# Allow sudo for wheel
sed -i 's/^# %wheel ALL=(ALL:ALL) ALL/%wheel ALL=(ALL:ALL) ALL/' /etc/sudoers

# 5. Install SonicDE (NATIVE PACMAN)
echo "------------------------------------------------"
echo "Installing SonicDE from Artix Repos..."
echo "------------------------------------------------"
# We assume sonicde-meta and sddm-openrc are in the enabled repos.
pacman -S --noconfirm sonicde-meta sddm-openrc

# 6. Bootloader
echo "Setting up Bootloader: $BOOTLOADER"
case "$BOOTLOADER" in
    "GRUB")
        pacman -S --noconfirm grub efibootmgr
        grub-install --target=x86_64-efi --efi-directory=/boot --bootloader-id=GRUB
        grub-mkconfig -o /boot/grub/grub.cfg
        ;;
    "rEFInd")
        pacman -S --noconfirm refind
        refind-install
        ;;
    "Limine")
        pacman -S --noconfirm limine
        mkdir -p /boot/EFI/BOOT
        cp /usr/share/limine/BOOTX64.EFI /boot/EFI/BOOT/BOOTX64.EFI
        ROOT_UUID=\$(findmnt / -n -o UUID)
        cat <<LIMINE > /boot/limine.conf
timeout: 5
/Artix Linux
    protocol: linux
    kernel_path: boot://2/vmlinuz-linux
    kernel_cmdline: root=UUID=\$ROOT_UUID rw rootflags=subvol=@
    module_path: boot://2/initramfs-linux.img
/Artix Linux (Fallback)
    protocol: linux
    kernel_path: boot://2/vmlinuz-linux
    kernel_cmdline: root=UUID=\$ROOT_UUID rw rootflags=subvol=@
    module_path: boot://2/initramfs-linux-fallback.img
LIMINE
        ;;
esac

# 7. Services
rc-update add sddm default
rc-update add elogind boot

echo "Configuration Complete."
EOF

    chmod +x "$MOUNT_POINT/setup_internal.sh"
    clear
    log "Entering Artix Chroot..."
    artix-chroot "$MOUNT_POINT" /setup_internal.sh
    if [ $? -ne 0 ]; then err "Chroot failed! Check errors."; fi
    rm "$MOUNT_POINT/setup_internal.sh"
}

# -- Main Execution --
pre_flight
setup_artix_repos "" 
select_menu_options
partition_drive
install_base
setup_artix_repos "$MOUNT_POINT"
configure_system

dialog --msgbox "Artix Installation Complete! Rebooting..." 10 60
umount -R "$MOUNT_POINT"
reboot