#! /bin/bash
#
# Copyright 2025, 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.
#

set -euo pipefail

mountedPoints=()
loopDev=""
mountRoot=""
rootPart=""
bootPart=""
imagePath=""

cleanup() {
    statusCode=$?
    set +e

    if [[ ${#mountedPoints[@]} -gt 0 ]]; then
        for ((index=${#mountedPoints[@]}-1; index>=0; index--)); do
            if mountpoint -q "${mountedPoints[$index]}"; then
                umount "${mountedPoints[$index]}" || true
            fi
        done
    fi

    if [[ -n "${mountRoot}" && -d "${mountRoot}" && "${mountRoot}" != "/mnt" ]]; then
        rmdir "${mountRoot}" 2>/dev/null || true
    fi

    if [[ -n "${loopDev}" ]]; then
        partx -d "${loopDev}" 2>/dev/null || true
        losetup --detach "${loopDev}" 2>/dev/null || true
    fi

    exit "$statusCode"
}

trap cleanup EXIT INT TERM

die() {
    echo "Error: $*" >&2
    exit 1
}

help() {
    echo "Usage: $0 <image-file>"
    echo "Mounts a Stormux image, chroots into it, and cleans up on exit."
    exit 0
}

require_cmd() {
    command -v "$1" >/dev/null 2>&1 || die "Missing required command: $1"
}

if [[ $# -ne 1 ]]; then
    help
fi

if [[ "$1" == "-h" || "$1" == "--help" ]]; then
    help
fi

imagePath="$(readlink -f "$1" 2>/dev/null || true)"
if [[ -z "${imagePath}" ]]; then
    die "Unable to resolve image path: $1"
fi
if [[ ! -f "${imagePath}" ]]; then
    die "Image file not found: ${imagePath}"
fi

if [[ "$(whoami)" != "root" ]]; then
    die "This script must be run as root."
fi

if command -v pacman >/dev/null 2>&1; then
    for packageName in arch-install-scripts util-linux; do
        if ! pacman -Q "$packageName" &>/dev/null; then
            die "Please install ${packageName} before continuing."
        fi
    done
else
    require_cmd arch-chroot
    require_cmd losetup
    require_cmd partx
    require_cmd lsblk
    require_cmd mount
    require_cmd umount
    require_cmd findmnt
    require_cmd mountpoint
fi

if [[ "$(uname -m)" == "x86_64" ]]; then
    if command -v pacman >/dev/null 2>&1; then
        if ! pacman -Q qemu-user-static &>/dev/null; then
            die "Please install qemu-user-static and qemu-user-static-binfmt before continuing."
        fi
        if ! pacman -Q qemu-user-static-binfmt &>/dev/null; then
            die "Please install qemu-user-static and qemu-user-static-binfmt before continuing."
        fi
    else
        require_cmd qemu-aarch64-static
    fi
fi

require_cmd arch-chroot
require_cmd losetup
require_cmd partx
require_cmd lsblk
require_cmd mount
require_cmd umount
require_cmd findmnt
require_cmd mountpoint

baseMountDir="/mnt"
if [[ ! -d "${baseMountDir}" || ! -w "${baseMountDir}" ]]; then
    die "/mnt is missing or not writable. Please create or fix permissions."
fi

if mountpoint -q "${baseMountDir}"; then
    if [[ -t 0 ]]; then
        echo "/mnt is already mounted. Unmount it to continue? [y/N]"
        read -r answerText
        case "${answerText}" in
            y|Y|yes|YES)
                umount "${baseMountDir}" || die "Failed to unmount /mnt"
                ;;
            *)
                die "Refusing to proceed while /mnt is mounted."
                ;;
        esac
    else
        die "/mnt is already mounted. Refusing to proceed in non-interactive mode."
    fi
fi

mountRoot="${baseMountDir}"

loopDev="$(losetup --find --show --partscan "${imagePath}")"
partx -u "${loopDev}" >/dev/null 2>&1 || true

declare -a rootCandidates=()
declare -a bootCandidates=()

while read -r deviceName fsType nodeType; do
    if [[ "${nodeType}" != "part" ]]; then
        continue
    fi
    case "${fsType}" in
        ext4|ext3|ext2)
            rootCandidates+=("${deviceName}")
            ;;
        vfat|fat|fat32|fat16)
            bootCandidates+=("${deviceName}")
            ;;
    esac
done < <(lsblk -nrpo NAME,FSTYPE,TYPE "${loopDev}")

if [[ ${#rootCandidates[@]} -gt 0 ]]; then
    rootPart="${rootCandidates[0]}"
elif [[ -b "${loopDev}p2" ]]; then
    rootPart="${loopDev}p2"
else
    die "Unable to locate root partition inside ${imagePath}"
fi

if [[ ${#bootCandidates[@]} -gt 0 ]]; then
    bootPart="${bootCandidates[0]}"
elif [[ -b "${loopDev}p1" ]]; then
    bootPart="${loopDev}p1"
fi

mount "${rootPart}" "${mountRoot}"
mountedPoints+=("${mountRoot}")

if [[ -n "${bootPart}" ]]; then
    mkdir -p "${mountRoot}/boot"
    mount "${bootPart}" "${mountRoot}/boot"
    mountedPoints+=("${mountRoot}/boot")
fi

echo "Mounted image at ${mountRoot}. Entering chroot..."
PS1="(Chroot) [\\u@\\h \\W] \\$ " arch-chroot "${mountRoot}"