#!/usr/bin/env bash
# synex-chroot - Simple chroot helper for Synex live systems
# Copyright (C) 2026 Synex Team <root@synex.ar>
# License: GPL-3.0-or-later

set -euo pipefail

VERSION="1.0.1"
TARGET="/mnt"
STATE_DIR="/run/synex-chroot"
RESOLV_BACKUP="${STATE_DIR}/resolv.conf.backup"
RESOLV_FLAG="${STATE_DIR}/resolv.conf.changed"

ROOT_DEV=""
BOOT_DEV=""
EFI_DEV=""
ROOT_FSTYPE=""
BTRFS_SUBVOL=""
BTRFS_ROOT_SUBVOL=""
MODE="oneshot"

CLEANUP_ON_ERROR="false"
CLEANUP_RUNNING="false"

Info ()
{
	printf 'synex-chroot: %s\n' "$*"
}

Cleanup_partial ()
{
	if [ "${CLEANUP_RUNNING}" = "true" ]; then
		return 0
	fi

	CLEANUP_RUNNING="true"

	set +e

	if command -v mountpoint >/dev/null 2>&1; then
		Restore_dns 2>/dev/null || true

		Umount_if_mounted "${TARGET}/boot/efi" 2>/dev/null || true
		Umount_if_mounted "${TARGET}/boot" 2>/dev/null || true
		Umount_if_mounted "${TARGET}/.snapshots" 2>/dev/null || true
		Umount_if_mounted "${TARGET}/var/log" 2>/dev/null || true
		Umount_if_mounted "${TARGET}/home" 2>/dev/null || true
		Umount_if_mounted "${TARGET}/run" 2>/dev/null || true
		Umount_if_mounted "${TARGET}/sys" 2>/dev/null || true
		Umount_if_mounted "${TARGET}/proc" 2>/dev/null || true
		Umount_if_mounted "${TARGET}/dev" 2>/dev/null || true
		Umount_if_mounted "${TARGET}" 2>/dev/null || true
	fi

	rmdir "${STATE_DIR}" 2>/dev/null || true

	set -e
}

Die ()
{
	printf 'synex-chroot: error: %s\n' "$*" >&2

	if [ "${CLEANUP_ON_ERROR}" = "true" ]; then
		Info "cleaning up after failed mount"
		Cleanup_partial
	fi

	exit 1
}

Usage ()
{
	cat <<'USAGE'
synex-chroot - Synex rescue chroot helper

Usage:
  synex-chroot ROOT_DEVICE [--boot BOOT_DEVICE] [--efi EFI_DEVICE]
  synex-chroot ROOT_DEVICE [--btrfs-subvol SUBVOL] [--boot BOOT_DEVICE] [--efi EFI_DEVICE]
  synex-chroot --mount ROOT_DEVICE [--boot BOOT_DEVICE] [--efi EFI_DEVICE]
  synex-chroot --mount ROOT_DEVICE [--btrfs-subvol SUBVOL] [--boot BOOT_DEVICE] [--efi EFI_DEVICE]
  synex-chroot --enter
  synex-chroot --umount
  synex-chroot --list
  synex-chroot --status
  synex-chroot --help

Description:
  synex-chroot helps prepare an installed Synex/Debian system under /mnt
  from a Synex live session. It mounts the target root filesystem, optional
  separate /boot and EFI partitions, binds the required runtime filesystems,
  and then enters the target system using chroot.

Supported target layouts:
  - Simple root filesystem: ext4, xfs, btrfs
  - Synex default Btrfs layout:
      @          -> /
      @home      -> /home
      @log       -> /var/log
      @snapshots -> /.snapshots
  - Optional separate /boot: ext4, xfs, btrfs
  - Optional separate EFI system partition: vfat/fat32

Btrfs behavior:
  If ROOT_DEVICE is Btrfs and /etc/os-release is not found after a direct
  mount, synex-chroot tries to detect the Synex default root subvolume "@"
  automatically. If found, it remounts ROOT_DEVICE using subvol=@.

  If @home, @log or @snapshots exist, they are mounted automatically on:
    /mnt/home
    /mnt/var/log
    /mnt/.snapshots

  A root subvolume can also be selected manually:
    --btrfs-subvol @

Unsupported by design:
  - LUKS encrypted devices
  - LVM activation
  - RAID assembly
  - ZFS root
  - Arbitrary complex Btrfs layouts beyond an explicit root subvolume
  - Automatic GRUB/initramfs/APT repair

Examples:
  List disks and filesystems:
    sudo synex-chroot --list

  One-shot mode: mount, enter chroot, and unmount when leaving:
    sudo synex-chroot /dev/nvme0n1p2

  One-shot mode with separate EFI partition:
    sudo synex-chroot /dev/nvme0n1p2 --efi /dev/nvme0n1p1

  One-shot mode with separate /boot and EFI partitions:
    sudo synex-chroot /dev/nvme0n1p2 --boot /dev/nvme0n1p3 --efi /dev/nvme0n1p1

  Synex default Btrfs layout, auto-detected:
    sudo synex-chroot /dev/nvme0n1p3 --boot /dev/nvme0n1p2 --efi /dev/nvme0n1p1

  Synex default Btrfs layout, explicit root subvolume:
    sudo synex-chroot /dev/nvme0n1p3 --btrfs-subvol @ --boot /dev/nvme0n1p2 --efi /dev/nvme0n1p1

  Mount only, without entering the chroot:
    sudo synex-chroot --mount /dev/nvme0n1p2 --efi /dev/nvme0n1p1

  Enter an already mounted target system:
    sudo synex-chroot --enter

  Unmount the prepared target system:
    sudo synex-chroot --umount

Typical repairs to run inside the chroot:
  update-initramfs -u -k all
  update-grub
  grub-install /dev/nvme0n1
  apt update
  apt install --reinstall linux-image-amd64 grub-efi-amd64-signed

Notes:
  - The target root is always mounted at /mnt.
  - This tool does not partition, format, decrypt, repair, or reinstall by itself.
  - All repair commands must be run manually inside the chroot.
USAGE
}

Require_root ()
{
	if [ "${EUID}" -ne 0 ]; then
		Die "this command must be run as root"
	fi
}

Have_command ()
{
	command -v "$1" >/dev/null 2>&1
}

Require_command ()
{
	Have_command "$1" || Die "required command not found: $1"
}

Canonical_device ()
{
	local DEV="$1"

	if [ ! -b "${DEV}" ]; then
		Die "not a block device: ${DEV}"
	fi

	readlink -f "${DEV}"
}

Device_fstype ()
{
	local DEV="$1"

	lsblk -no FSTYPE "${DEV}" 2>/dev/null | head -n 1 | tr '[:upper:]' '[:lower:]'
}

Validate_root_fs ()
{
	local DEV="$1"
	local FSTYPE

	FSTYPE="$(Device_fstype "${DEV}")"

	case "${FSTYPE}" in
		ext4|xfs|btrfs)
			ROOT_FSTYPE="${FSTYPE}"
			return 0
			;;

		"")
			Die "could not detect filesystem type for ${DEV}"
			;;

		*)
			Die "unsupported root filesystem on ${DEV}: ${FSTYPE}"
			;;
	esac
}

Validate_boot_fs ()
{
	local DEV="$1"
	local FSTYPE

	FSTYPE="$(Device_fstype "${DEV}")"

	case "${FSTYPE}" in
		ext4|xfs|btrfs)
			return 0
			;;

		"")
			Die "could not detect filesystem type for ${DEV}"
			;;

		*)
			Die "unsupported /boot filesystem on ${DEV}: ${FSTYPE}"
			;;
	esac
}

Validate_efi_fs ()
{
	local DEV="$1"
	local FSTYPE

	FSTYPE="$(Device_fstype "${DEV}")"

	case "${FSTYPE}" in
		vfat|fat|fat32)
			return 0
			;;

		"")
			Die "could not detect filesystem type for ${DEV}"
			;;

		*)
			Die "unsupported EFI filesystem on ${DEV}: ${FSTYPE}"
			;;
	esac
}

Is_mounted ()
{
	mountpoint -q "$1"
}

Looks_like_target_system ()
{
	[ -d "${TARGET}/etc" ] && [ -f "${TARGET}/etc/os-release" ]
}

Mountpoint_busy_check ()
{
	if Is_mounted "${TARGET}"; then
		Die "${TARGET} is already mounted; run 'synex-chroot --status' or 'synex-chroot --umount' first"
	fi
}

Ensure_target_system ()
{
	if ! Looks_like_target_system; then
		Die "${TARGET} does not look like a valid installed Linux system"
	fi
}

Mount_device ()
{
	local DEV="$1"
	local DIR="$2"
	local OPTIONS="${3:-}"

	mkdir -p "${DIR}"

	if [ -n "${OPTIONS}" ]; then
		Info "mounting ${DEV} on ${DIR} with options: ${OPTIONS}"
		mount -o "${OPTIONS}" "${DEV}" "${DIR}"
	else
		Info "mounting ${DEV} on ${DIR}"
		mount "${DEV}" "${DIR}"
	fi
}

Detect_btrfs_root_subvol_from_top_level ()
{
	local SUBVOL

	if [ -n "${BTRFS_SUBVOL}" ]; then
		printf '%s\n' "${BTRFS_SUBVOL}"
		return 0
	fi

	for SUBVOL in @ @root root rootfs; do
		if [ -f "${TARGET}/${SUBVOL}/etc/os-release" ]; then
			printf '%s\n' "${SUBVOL}"
			return 0
		fi
	done

	return 1
}

Remount_btrfs_root_subvol ()
{
	local SUBVOL

	SUBVOL="$(Detect_btrfs_root_subvol_from_top_level)" || {
		Die "Btrfs filesystem detected, but no valid root subvolume was found; try '--btrfs-subvol @'"
	}

	Info "detected Btrfs root subvolume: ${SUBVOL}"

	Umount_if_mounted "${TARGET}"

	Mount_device "${ROOT_DEV}" "${TARGET}" "subvol=${SUBVOL}"

	BTRFS_ROOT_SUBVOL="${SUBVOL}"
}

Btrfs_subvol_exists ()
{
	local SUBVOL="$1"
	local TMP_DIR
	local RET

	RET="1"
	TMP_DIR="$(mktemp -d /run/synex-chroot-btrfs.XXXXXX)"

	if mount -o subvolid=5 "${ROOT_DEV}" "${TMP_DIR}" >/dev/null 2>&1; then
		if [ -d "${TMP_DIR}/${SUBVOL}" ]; then
			RET="0"
		fi

		umount "${TMP_DIR}" >/dev/null 2>&1 || true
	fi

	rmdir "${TMP_DIR}" >/dev/null 2>&1 || true

	return "${RET}"
}

Mount_btrfs_subvol_if_exists ()
{
	local SUBVOL="$1"
	local DEST="$2"

	if Btrfs_subvol_exists "${SUBVOL}"; then
		mkdir -p "${TARGET}/${DEST}"

		if ! Is_mounted "${TARGET}/${DEST}"; then
			Info "mounting Btrfs subvolume ${SUBVOL} on ${TARGET}/${DEST}"
			mount -o "subvol=${SUBVOL}" "${ROOT_DEV}" "${TARGET}/${DEST}"
		fi
	fi
}

Mount_synex_btrfs_subvolumes ()
{
	if [ "${ROOT_FSTYPE}" != "btrfs" ]; then
		return 0
	fi

	if [ "${BTRFS_ROOT_SUBVOL}" != "@" ]; then
		return 0
	fi

	Mount_btrfs_subvol_if_exists "@home" "home"
	Mount_btrfs_subvol_if_exists "@log" "var/log"
	Mount_btrfs_subvol_if_exists "@snapshots" ".snapshots"
}

Bind_mount ()
{
	local SOURCE="$1"
	local DEST="$2"

	mkdir -p "${DEST}"

	if ! Is_mounted "${DEST}"; then
		Info "binding ${SOURCE} on ${DEST}"
		mount --rbind "${SOURCE}" "${DEST}"
		mount --make-rslave "${DEST}"
	fi
}

Mount_api_filesystems ()
{
	Bind_mount /dev "${TARGET}/dev"

	mkdir -p "${TARGET}/proc"
	if ! Is_mounted "${TARGET}/proc"; then
		Info "mounting proc on ${TARGET}/proc"
		mount -t proc proc "${TARGET}/proc"
	fi

	Bind_mount /sys "${TARGET}/sys"
	Bind_mount /run "${TARGET}/run"
}

Prepare_dns ()
{
	mkdir -p "${STATE_DIR}"

	if [ ! -d "${TARGET}/etc" ]; then
		return 0
	fi

	if [ -L "${TARGET}/etc/resolv.conf" ]; then
		Info "leaving target resolv.conf symlink unchanged"
		return 0
	fi

	if [ -f /etc/resolv.conf ]; then
		if [ -f "${TARGET}/etc/resolv.conf" ] && [ ! -f "${RESOLV_BACKUP}" ]; then
			cp -a "${TARGET}/etc/resolv.conf" "${RESOLV_BACKUP}"
		fi

		cp -L /etc/resolv.conf "${TARGET}/etc/resolv.conf"
		touch "${RESOLV_FLAG}"
		Info "copied host DNS resolver configuration into target"
	fi
}

Restore_dns ()
{
	if [ -f "${RESOLV_FLAG}" ]; then
		if [ -f "${RESOLV_BACKUP}" ]; then
			cp -a "${RESOLV_BACKUP}" "${TARGET}/etc/resolv.conf" 2>/dev/null || true
			Info "restored target DNS resolver configuration"
		fi

		rm -f "${RESOLV_FLAG}" "${RESOLV_BACKUP}"
	fi
}

Mount_target ()
{
	Require_root
	Require_command lsblk
	Require_command mount
	Require_command mountpoint
	Require_command chroot
	Require_command mktemp

	CLEANUP_ON_ERROR="true"

	ROOT_DEV="$(Canonical_device "${ROOT_DEV}")"
	Validate_root_fs "${ROOT_DEV}"

	if [ -n "${BOOT_DEV}" ]; then
		BOOT_DEV="$(Canonical_device "${BOOT_DEV}")"
		Validate_boot_fs "${BOOT_DEV}"
	fi

	if [ -n "${EFI_DEV}" ]; then
		EFI_DEV="$(Canonical_device "${EFI_DEV}")"
		Validate_efi_fs "${EFI_DEV}"
	fi

	if [ -n "${BTRFS_SUBVOL}" ] && [ "${ROOT_FSTYPE}" != "btrfs" ]; then
		Die "--btrfs-subvol can only be used with a Btrfs root filesystem"
	fi

	Mountpoint_busy_check
	mkdir -p "${TARGET}"

	if [ "${ROOT_FSTYPE}" = "btrfs" ] && [ -n "${BTRFS_SUBVOL}" ]; then
		Mount_device "${ROOT_DEV}" "${TARGET}" "subvol=${BTRFS_SUBVOL}"
		BTRFS_ROOT_SUBVOL="${BTRFS_SUBVOL}"
	else
		Mount_device "${ROOT_DEV}" "${TARGET}"
	fi

	if ! Looks_like_target_system; then
		if [ "${ROOT_FSTYPE}" = "btrfs" ] && [ -z "${BTRFS_ROOT_SUBVOL}" ]; then
			Remount_btrfs_root_subvol
		fi
	fi

	Ensure_target_system
	Mount_synex_btrfs_subvolumes

	if [ -n "${BOOT_DEV}" ]; then
		Mount_device "${BOOT_DEV}" "${TARGET}/boot"
	fi

	if [ -n "${EFI_DEV}" ]; then
		mkdir -p "${TARGET}/boot/efi"
		Mount_device "${EFI_DEV}" "${TARGET}/boot/efi"
	fi

	Mount_api_filesystems
	Prepare_dns

	CLEANUP_ON_ERROR="false"

	Info "target system is ready under ${TARGET}"
}

Enter_target ()
{
	Require_root
	Require_command chroot
	Require_command mountpoint

	if ! Is_mounted "${TARGET}"; then
		Die "${TARGET} is not mounted; use 'synex-chroot ROOT_DEVICE' or 'synex-chroot --mount ROOT_DEVICE' first"
	fi

	Ensure_target_system
	Mount_api_filesystems
	Prepare_dns

	Info "entering chroot; type 'exit' to leave"

	SHELL_PATH="/bin/bash"
	SHELL_ARGS=("-l")

	if [ ! -x "${TARGET}${SHELL_PATH}" ]; then
		SHELL_PATH="/bin/sh"
		SHELL_ARGS=()
	fi

	chroot "${TARGET}" /usr/bin/env -i \
		HOME=/root \
		TERM="${TERM:-linux}" \
		PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin \
		SYNEX_CHROOT=1 \
		"${SHELL_PATH}" "${SHELL_ARGS[@]}"
}

Umount_if_mounted ()
{
	local DIR="$1"

	if Is_mounted "${DIR}"; then
		Info "unmounting ${DIR}"
		umount -R "${DIR}" 2>/dev/null || umount "${DIR}"
	fi
}

Umount_target ()
{
	Require_root
	Require_command umount
	Require_command mountpoint

	Restore_dns

	Umount_if_mounted "${TARGET}/boot/efi"
	Umount_if_mounted "${TARGET}/boot"
	Umount_if_mounted "${TARGET}/.snapshots"
	Umount_if_mounted "${TARGET}/var/log"
	Umount_if_mounted "${TARGET}/home"
	Umount_if_mounted "${TARGET}/run"
	Umount_if_mounted "${TARGET}/sys"
	Umount_if_mounted "${TARGET}/proc"
	Umount_if_mounted "${TARGET}/dev"
	Umount_if_mounted "${TARGET}"

	rmdir "${STATE_DIR}" 2>/dev/null || true
	Info "done"
}

List_devices ()
{
	Require_command lsblk
	lsblk -o NAME,TYPE,FSTYPE,LABEL,UUID,SIZE,MOUNTPOINTS
}

Status ()
{
	Require_command findmnt
	Require_command mountpoint

	printf 'synex-chroot status\n'
	printf 'target: %s\n\n' "${TARGET}"

	if Is_mounted "${TARGET}"; then
		findmnt -R "${TARGET}"
	else
		printf '%s is not mounted\n' "${TARGET}"
	fi
}

Parse_args ()
{
	if [ "$#" -eq 0 ]; then
		Usage
		exit 1
	fi

	case "$1" in
		--help|-h)
			Usage
			exit 0
			;;

		--version|-V)
			printf 'synex-chroot %s\n' "${VERSION}"
			exit 0
			;;

		--list)
			MODE="list"
			shift
			;;

		--status)
			MODE="status"
			shift
			;;

		--enter)
			MODE="enter"
			shift
			;;

		--umount|--unmount)
			MODE="umount"
			shift
			;;

		--mount)
			MODE="mount"
			shift
			[ "$#" -gt 0 ] || Die "--mount requires ROOT_DEVICE"
			ROOT_DEV="$1"
			shift
			;;

		-*)
			Die "unknown option: $1"
			;;

		*)
			MODE="oneshot"
			ROOT_DEV="$1"
			shift
			;;
	esac

	while [ "$#" -gt 0 ]; do
		case "$1" in
			--boot)
				shift
				[ "$#" -gt 0 ] || Die "--boot requires BOOT_DEVICE"
				BOOT_DEV="$1"
				shift
				;;

			--efi)
				shift
				[ "$#" -gt 0 ] || Die "--efi requires EFI_DEVICE"
				EFI_DEV="$1"
				shift
				;;

			--btrfs-subvol)
				shift
				[ "$#" -gt 0 ] || Die "--btrfs-subvol requires SUBVOL"
				BTRFS_SUBVOL="$1"
				shift
				;;

			*)
				Die "unknown argument: $1"
				;;
		esac
	done

	case "${MODE}" in
		mount|oneshot)
			[ -n "${ROOT_DEV}" ] || Die "root device is required"
			;;

		list|status|enter|umount)
			if [ -n "${ROOT_DEV}" ] || [ -n "${BOOT_DEV}" ] ||
			   [ -n "${EFI_DEV}" ] || [ -n "${BTRFS_SUBVOL}" ]; then
				Die "${MODE} does not accept device arguments"
			fi
			;;
	esac
}

Main ()
{
	Parse_args "$@"

	case "${MODE}" in
		list)
			List_devices
			;;

		status)
			Status
			;;

		mount)
			Mount_target
			;;

		enter)
			Enter_target
			;;

		umount)
			Umount_target
			;;

		oneshot)
			Mount_target
			set +e
			Enter_target
			RET="$?"
			set -e
			Umount_target
			exit "${RET}"
			;;

		*)
			Die "internal error: invalid mode ${MODE}"
			;;
	esac
}

Main "$@"
