#! /bin/bash
#
# Copyright 2020, Stormux, <storm_dragon@stormux.org>
#
# This is free software; you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free
# Software Foundation; either version 3, or (at your option) any later
# version.
#
# This software is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this package; see the file COPYING.  If not, write to the Free
# Software Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.
#

# keep track of mounted status for exit function
mounted=1

set -e # Don't want to destroy stuff if this goes majorly wrong.
trap cleanup EXIT # make sure the script cleans up after itself before closing.


# shellcheck disable=SC2329 # cleanup is invoked via trap EXIT
cleanup() {
    status=$? # capture original exit status so failures propagate

    if [[ $mounted -eq 0 ]]; then
        umount -R /mnt || true
    fi

    if [[ -n "${loopdev:-}" ]]; then
        partx -d "${loopdev}" || true
        losetup --detach "${loopdev}" || true
    fi

    # Clean up temporary pacman config directory
    if [[ -n "${tmpDir:-}" && -d "${tmpDir:-}" ]]; then
        rm -rf "${tmpDir}"
    fi

    exit "$status"
}

help() {
    echo -e "Usage:\n"
    echo "With no arguments, build with default parameters."
    for i in "${!command[@]}" ; do
        echo "-${i/:/ <parameter>}: ${command[${i}]}"
    done | sort
    exit 0
}

# Array of command line arguments
declare -A command=(
    [h]="This help screen."
    [l:]="Language default is en_US."
    [n:]="Image name, default is stormux-pi4-aarch64-<yyyy-mm-dd>.img"
    [s:]="image size in GB, default is 6."
)

# Convert the keys of the associative array to a format usable by getopts
args="${!command[*]}"
args="${args//[[:space:]]/}"
while getopts "${args}" i ; do
    case "$i" in
        h) help ;;
        l)
            imageLanguage="${OPTARG}.UTF-8"
            ;;
        n)
            imageName="${OPTARG}"
            ;;
        s)
            if [[ "${OPTARG}" =~ ^[[:digit:]]+$ ]]; then
                imageSize="${OPTARG}G"
            else
                echo "Image size must be numeric."
                exit 1
            fi
            ;;
        *)
            exit 1
            ;;
    esac
done

# make sure variables are set, or use defaults.
export imageSize="${imageSize:-6G}"
imageName="${imageName:-stormux-pi4-aarch64-$(date '+%Y-%m-%d').img}"
export imageName
export imageLanguage="${imageLanguage:-en_US.UTF-8}"

# Make sure the image file doesn't exist.
if [[ -e "$imageName" ]]; then
    echo "${imageName} exists, please remove or move it for this script to continue."
    exit 1
fi

# Make sure this script is ran as root.
if [ "$(whoami)" != "root" ] ; then
    echo "Error: This script must be run as root."
    exit 1
fi

# make sure the needed tools are installed
if [[ "$(uname -m)" == "x86_64" ]]; then
    if ! pacman -Q qemu-user-static &> /dev/null ; then
        echo "Please install qemu-user-static and qemu-user-static-binfmt before continuing."
        exit 1
    fi
    if ! pacman -Q qemu-user-static-binfmt &> /dev/null ; then
        echo "Please install qemu-user-static and qemu-user-static-binfmt before continuing."
        exit 1
    fi
fi
for i in arch-install-scripts dosfstools parted ; do
    if ! pacman -Q $i &> /dev/null ; then
        echo "Please install $i before continuing."
        exit 1
    fi
done


fallocate -l "$imageSize" "$imageName"
loopdev="$(losetup --find --show "${imageName}")"
parted --script "${loopdev}" mklabel msdos mkpart primary fat32 0% 512M mkpart primary ext4 512M 100%
mkfs.vfat -F32 -n STRMX_BOOT "${loopdev}p1"
mkfs.ext4 -F -L STRMX_ROOT "${loopdev}p2"
mount "${loopdev}p2" /mnt
mkdir /mnt/boot
mount "${loopdev}p1" /mnt/boot
# Things are mounted now, so set mounted to 0 (bash true)
mounted=0

# Set up temporary pacman configuration for ARM bootstrap
echo "Setting up pacman configuration for aarch64..."
tmpDir=$(mktemp -d)
mkdir -p "${tmpDir}/pacman.d"

# Create mirrorlist for Arch Linux ARM
cat > "${tmpDir}/pacman.d/mirrorlist" << 'MIRRORLIST'
Server = http://mirror.archlinuxarm.org/$arch/$repo
Server = http://fl.us.mirror.archlinuxarm.org/$arch/$repo
Server = http://il.us.mirror.archlinuxarm.org/$arch/$repo
Server = http://tx.us.mirror.archlinuxarm.org/$arch/$repo
Server = http://nj.us.mirror.archlinuxarm.org/$arch/$repo
MIRRORLIST

# Create pacman.conf for ARM
cat > "${tmpDir}/pacman.conf" << PACMANCONF
[options]
HoldPkg     = pacman glibc
Architecture = aarch64
CheckSpace
SigLevel    = Never
LocalFileSigLevel = Never

[core]
Include = ${tmpDir}/pacman.d/mirrorlist

[extra]
Include = ${tmpDir}/pacman.d/mirrorlist

[alarm]
Include = ${tmpDir}/pacman.d/mirrorlist
PACMANCONF

# Bootstrap the system with pacstrap
echo "Running pacstrap to install base system..."
pacstrap -c -G -M -C "${tmpDir}/pacman.conf" /mnt base base-devel linux-rpi raspberrypi-bootloader firmware-raspberrypi archlinuxarm-keyring

# Clean up temporary config
rm -rf "${tmpDir}"

# Copy override files into place before chroot (except skel to avoid conflicts)
cp -rv ../files/boot/* /mnt/boot
cp -rv ../files/var/* /mnt/var
cp -rv ../files/usr/* /mnt/usr
# Copy etc files but exclude skel directory to avoid package conflicts
find ../files/etc -mindepth 1 -maxdepth 1 ! -name skel -exec cp -rv {} /mnt/etc/ \;

PS1="(Chroot) [\u@\h \W] \$" arch-chroot /mnt << EOF
echo "Chroot started."
set -euo pipefail
# Disable pacman sandboxing (doesn't work in chroot)
sed -i '/^\[options\]/a DisableSandbox' /etc/pacman.conf
# set up pacman
pacman-key --init
pacman-key --populate archlinuxarm

# Add stormux repository
echo "Downloading stormux repository key..."
curl -s https://packages.stormux.org/stormux_repo.pub > /tmp/stormux_repo.pub
echo "Adding key to pacman keyring..."
pacman-key --add /tmp/stormux_repo.pub
echo "Locally signing key..."
pacman-key --lsign-key 52ADA49000F1FF0456F8AEEFB4CDE1CD56EF8E82
rm /tmp/stormux_repo.pub
# Add repository before core in pacman.conf
sed -i '/^\[core\]/i[stormux]\nSigLevel = Required DatabaseOptional\nServer = https://packages.stormux.org/\$arch\n' /etc/pacman.conf

# Test mirrors and configure based on responsiveness
echo "Testing mirror responsiveness..."
mirrors=(
    "mirror.archlinuxarm.org"
    "fl.us.mirror.archlinuxarm.org"
    "il.us.mirror.archlinuxarm.org"
    "tx.us.mirror.archlinuxarm.org"
    "nj.us.mirror.archlinuxarm.org"
)

working_mirrors=()
for mirror in "\${mirrors[@]}"; do
    echo "Testing \$mirror..."
    if curl -m 5 -f -s "http://\$mirror/aarch64/core/core.db" > /dev/null 2>&1; then
        echo "  \$mirror: OK"
        working_mirrors+=("\$mirror")
    else
        echo "  \$mirror: FAILED"
    fi
done

# Create mirrorlist with working mirrors first, then fallback to all mirrors
> /etc/pacman.d/mirrorlist
if [[ \${#working_mirrors[@]} -gt 0 ]]; then
    echo "Using \${#working_mirrors[@]} working mirrors"
    for mirror in "\${working_mirrors[@]}"; do
        echo "Server = http://\$mirror/\\\$arch/\\\$repo" >> /etc/pacman.d/mirrorlist
    done
else
    echo "No mirrors responded, using all mirrors as fallback"
    for mirror in "\${mirrors[@]}"; do
        echo "Server = http://\$mirror/\\\$arch/\\\$repo" >> /etc/pacman.d/mirrorlist
    done
fi

pacman -Syy

# Install all packages (stormux repo has priority since it's listed first)
packages=(
    alsa-firmware
    alsa-utils
    base
    base-devel
    bash-completion
    bluez
    bluez-utils
    brltty
    cronie
    dialog
    espeak-ng
    fake-hwclock
    fenrir
    firmware-raspberrypi
    linux-firmware
    git
    gstreamer
    gst-plugins-base
    gst-plugins-good
    ii
    magic-wormhole
    man
    man-pages
    networkmanager
    openssh
    parted
    pipewire
    pipewire-alsa
    pipewire-jack
    pipewire-pulse
    poppler
    python-dbus
    python-gobject
    python-pyenchant
    python-pyte
    python-pyperclip
    raspberrypi-utils
    socat
    realtime-privileges
    rhvoice-voice-bdl
    rng-tools
    rsync
    screen
    sox
    w3m-git
    wget
    wireless-regdb
    wireplumber
    vi
    xdg-user-dirs
    xdg-utils
    yay
)

pacman -Su --needed --noconfirm "\${packages[@]}"

# Fix mkinitcpio preset for linux-rpi - kms hook fails on aarch64
# See: https://archlinuxarm.org/forum/viewtopic.php?t=16672
if [[ -f /etc/mkinitcpio.d/linux-rpi.preset ]]; then
    echo "Configuring mkinitcpio preset to skip kms hook..."
    sed -i "s/^default_options=.*/default_options=\"--skiphook kms\"/" /etc/mkinitcpio.d/linux-rpi.preset
    sed -i "s/^fallback_options=.*/fallback_options=\"-S autodetect --skiphook kms\"/" /etc/mkinitcpio.d/linux-rpi.preset
fi

# Regenerate initramfs with the corrected settings
echo "Regenerating initramfs..."
mkinitcpio -P

# Restart gpg agents.
gpgconf --kill all
# set the language
sed -i "s/#$imageLanguage/$imageLanguage/" /etc/locale.gen
echo "LANG=$imageLanguage" > /etc/locale.conf
locale-gen
# Configure and enable Hardware Random Number Generator
echo 'RNGD_OPTS="-o /dev/random -r /dev/hwrng"' > /etc/conf.d/rngd
systemctl enable rngd.service
# Set the distribution name.
echo 'Stormux \r (\l)' > /etc/issue
echo >> /etc/issue
# Create the stormux user
useradd -m -g users -G wheel,realtime,audio,video,network,brlapi -s /bin/bash stormux
# Grant sudo privileges to the stormux user for package installation.
echo 'stormux ALL=(ALL) NOPASSWD: ALL' >> /etc/sudoers.d/wheel
# Set the password for the root user
echo -e "root\nroot" | passwd "root"
# Set the password for the stormux user
echo -e "stormux\nstormux" | passwd "stormux"

# Change to the stormux user for user configuration
sudo -iu stormux
# suppress git spam about default branch name
git config --global init.defaultBranch master
# Create desktop, downloads, music, and other directories.
xdg-user-dirs-update
exit
# Enable linger so that hopefully sound will start at login.
mkdir -p /var/lib/systemd/linger
touch /var/lib/systemd/linger/stormux
systemctl --global enable pipewire.service pipewire-pulse.service
/usr/share/fenrirscreenreader/tools/configure_pipewire.sh
sudo -u stormux /usr/share/fenrirscreenreader/tools/configure_pipewire.sh
# Configure sudo for group wheel, remove nopasswd for the stormux user
echo '%wheel ALL=(ALL) ALL' > /etc/sudoers.d/wheel
# Set the hostname
echo stormux > /etc/hostname
# Configure services
systemctl disable systemd-networkd.service systemd-networkd.socket
# Enable services
services=(
    brltty.path
    cronie.service
    fake-hwclock.service
    fenrirscreenreader.service
    log-to-ram-setup.service
    log-to-ram-sync.timer
    log-to-ram-shutdown.service
    NetworkManager.service
    sshd.service
    ssh-login-monitor.service
)

for service in "\${services[@]}"; do
    echo "Enabling \$service..."
    if systemctl enable "\$service"; then
        echo "  \$service: OK"
    else
        echo "  \$service: FAILED"
    fi
done

pacman -Sc --noconfirm

# Re-enable pacman sandboxing for the final image
sed -i '/^DisableSandbox/d' /etc/pacman.conf

EOF

# Copy skel files to stormux user home (after user rename in chroot)
find ../files/etc/skel/ -mindepth 1 -exec cp -rv "{}" /mnt/home/stormux/ \;

# Copy boot files again to ensure custom config overrides any package changes
cp -rv ../files/boot/* /mnt/boot

# Exiting calls the cleanup function to unmount.
exit 0