#!/usr/bin/env bash
# SPDX-FileCopyrightText: Daniel Pätzold
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# Bootstrap script: creates an OEMDRV partition by shrinking an existing
# ext4 or btrfs partition by 4 GiB, formats the freed space as BTRFS with
# label OEMDRV, mounts it to /opt/sys_config and clones the repository.
#
# Run as root on a target machine before any Kickstart installation.

SHRINK_MIB=4096
OEMDRV_LABEL="OEMDRV"
MOUNT_POINT="/opt/sys_config"
MOUNT_OPTS="compress=zstd:6"
REPO_URL="https://gitea.dtext.online/obel1x/fedora-OEMDRV.git"
MIN_FREE_MIB=$(( SHRINK_MIB + 512 ))   # require 512 MiB headroom above the shrink size

# ── Helpers ───────────────────────────────────────────────────────────────────

die()  { echo; echo "ERROR: $*" >&2; exit 1; }
info() { echo; echo ">>> $*"; }
hr()   { printf '%.0s─' {1..100}; echo; }

require_root() {
    [[ "$EUID" -eq 0 ]] || die "This script must be run as root."
}

check_tools() {
    local missing=()
    for tool in lsblk blkid parted partprobe mkfs.btrfs git e2fsck resize2fs tune2fs; do
        command -v "$tool" >/dev/null 2>&1 || missing+=("$tool")
    done
    [[ ${#missing[@]} -eq 0 ]] || die "Missing required tools: ${missing[*]}"
}

# ── Free-space helpers ────────────────────────────────────────────────────────

# Free MiB for a mounted device via df
mounted_free_mib() {
    df --output=avail -BM "$1" 2>/dev/null | awk 'NR==2{gsub("M",""); print $1+0}'
}

# Free MiB for an unmounted ext4 device via tune2fs (no mount needed)
ext4_free_mib() {
    local dev="$1" fb bs
    fb=$(tune2fs -l "$dev" 2>/dev/null | awk '/^Free blocks:/{print $3}')
    bs=$(tune2fs -l "$dev" 2>/dev/null | awk '/^Block size:/{print $3}')
    [[ -n "$fb" && -n "$bs" ]] || { echo 0; return; }
    echo $(( fb * bs / 1048576 ))
}

# Free MiB for a btrfs device – mounts temporarily if not already mounted
btrfs_free_mib() {
    local dev="$1" mnt free tmp=0
    mnt=$(lsblk -n -o MOUNTPOINT "$dev" 2>/dev/null | grep -v '^$' | head -1)
    if [[ -z "$mnt" ]]; then
        mnt=$(mktemp -d /tmp/oemdrv_check.XXXXXX)
        mount -o ro "$dev" "$mnt" 2>/dev/null && tmp=1 || { rmdir "$mnt"; echo 0; return; }
    fi
    free=$(mounted_free_mib "$mnt")
    [[ $tmp -eq 1 ]] && { umount "$mnt"; rmdir "$mnt"; }
    echo "${free:-0}"
}

# ── Partition discovery ───────────────────────────────────────────────────────

# Parallel arrays, indexed 0..PT_IDX-1
PT_IDX=0
PT_DEV=()   # /dev/sda2
PT_DISK=()  # /dev/sda
PT_PNUM=()  # 2
PT_SIZE=()  # "100G"
PT_FS=()    # ext4 | btrfs
PT_LABEL=() # home | -
PT_UUID=()  # xxxxxxxx | -
PT_MNT=()   # /home | -
PT_FREE=()  # free MiB (integer)
PT_OK=()    # 1 = shrinkable

collect_partitions() {
    local NAME SIZE FSTYPE LABEL UUID MOUNTPOINT TYPE PKNAME

    while IFS= read -r line; do
        NAME="" SIZE="" FSTYPE="" LABEL="" UUID="" MOUNTPOINT="" TYPE="" PKNAME=""
        eval "$line"

        [[ "$TYPE"   != "part"                         ]] && continue
        [[ "$FSTYPE" != "ext4" && "$FSTYPE" != "btrfs" ]] && continue
        [[ "$LABEL"  == "$OEMDRV_LABEL"                ]] && continue
        [[ -z "$PKNAME"                                ]] && continue

        local disk="/dev/$PKNAME"
        local pnum
        pnum=$(cat /sys/class/block/"${NAME##/dev/}"/partition 2>/dev/null) || continue

        # Free space check
        local free_mib
        if [[ -n "$MOUNTPOINT" ]]; then
            free_mib=$(mounted_free_mib "$NAME")
        elif [[ "$FSTYPE" == "ext4" ]]; then
            free_mib=$(ext4_free_mib "$NAME")
        else
            free_mib=$(btrfs_free_mib "$NAME")
        fi
        free_mib=${free_mib:-0}

        # Shrinkable: not /, enough free space
        local ok=0
        [[ "$MOUNTPOINT" != "/" && "$free_mib" -ge "$MIN_FREE_MIB" ]] && ok=1

        PT_DEV[$PT_IDX]="$NAME"
        PT_DISK[$PT_IDX]="$disk"
        PT_PNUM[$PT_IDX]="$pnum"
        PT_SIZE[$PT_IDX]="$SIZE"
        PT_FS[$PT_IDX]="$FSTYPE"
        PT_LABEL[$PT_IDX]="${LABEL:--}"
        PT_UUID[$PT_IDX]="${UUID:--}"
        PT_MNT[$PT_IDX]="${MOUNTPOINT:--}"
        PT_FREE[$PT_IDX]="$free_mib"
        PT_OK[$PT_IDX]="$ok"
        (( PT_IDX++ )) || true
    done < <(lsblk -o NAME,SIZE,FSTYPE,LABEL,UUID,MOUNTPOINT,TYPE,PKNAME -p -P -n 2>/dev/null)
}

# ── Table display ─────────────────────────────────────────────────────────────

show_table() {
    hr
    printf "  %-4s  %-14s  %-6s  %-8s  %-10s  %-6s  %-18s  %-36s  %s\n" \
        "#" "Device" "Disk" "Size" "Free (MiB)" "FS" "Label" "UUID" "Mountpoint"
    hr
    local i mark note
    for (( i=0; i<PT_IDX; i++ )); do
        mark="  "
        note=""
        if [[ "${PT_OK[$i]}" == "1" ]]; then
            mark="* "
        elif [[ "${PT_MNT[$i]}" == "/" ]]; then
            note="  (root — skip)"
        else
            note="  (not enough free space)"
        fi
        printf "%s%-4s  %-14s  %-6s  %-8s  %-10s  %-6s  %-18s  %-36s  %s%s\n" \
            "$mark" \
            "$((i+1))" \
            "${PT_DEV[$i]}" \
            "${PT_DISK[$i]##/dev/}" \
            "${PT_SIZE[$i]}" \
            "${PT_FREE[$i]}" \
            "${PT_FS[$i]}" \
            "${PT_LABEL[$i]}" \
            "${PT_UUID[$i]}" \
            "${PT_MNT[$i]}" \
            "$note"
    done
    hr
    printf "  *  = shrinkable: will be reduced by %s MiB to create OEMDRV\n" "$SHRINK_MIB"
    echo
}

# ── Partition geometry via parted -m ─────────────────────────────────────────

part_end_mib() {
    parted -s "$1" -m unit MiB print 2>/dev/null \
        | awk -F: -v p="$2" '$1==p{gsub("MiB","",$3); printf "%.0f\n",$3; exit}'
}

part_start_mib() {
    parted -s "$1" -m unit MiB print 2>/dev/null \
        | awk -F: -v p="$2" '$1==p{gsub("MiB","",$2); printf "%.0f\n",$2; exit}'
}

# Device name for a new partition on a given disk
new_part_device() {
    local disk="$1" pnum="$2"
    if [[ "$disk" =~ (nvme|mmcblk) ]]; then
        echo "${disk}p${pnum}"
    else
        echo "${disk}${pnum}"
    fi
}

# ── Main ──────────────────────────────────────────────────────────────────────

require_root
check_tools

info "Scanning partitions for shrinkable ext4 / btrfs volumes..."
collect_partitions

[[ $PT_IDX -gt 0 ]] || die "No ext4 or btrfs partitions found on this system."

show_table

# Count shrinkable
shrink_count=0
for (( i=0; i<PT_IDX; i++ )); do [[ "${PT_OK[$i]}" == "1" ]] && (( shrink_count++ )) || true; done
[[ $shrink_count -gt 0 ]] || die "No shrinkable partitions found. A non-root ext4 or btrfs partition needs at least ${MIN_FREE_MIB} MiB free."

# ── Selection ─────────────────────────────────────────────────────────────────

SEL=-1
while true; do
    read -r -p "Enter number of partition to shrink, or q to quit: " INPUT
    [[ "$INPUT" =~ ^[qQ]$ ]] && { echo "Aborted."; exit 0; }
    [[ "$INPUT" =~ ^[0-9]+$ ]] || { echo "  Please enter a number."; continue; }
    local_i=$(( INPUT - 1 ))
    (( local_i >= 0 && local_i < PT_IDX )) || { echo "  Number out of range."; continue; }
    [[ "${PT_OK[$local_i]}" == "1" ]] || { echo "  That partition cannot be shrunk (see marks above)."; continue; }
    SEL=$local_i
    break
done

S_DEV="${PT_DEV[$SEL]}"
S_DISK="${PT_DISK[$SEL]}"
S_PNUM="${PT_PNUM[$SEL]}"
S_FS="${PT_FS[$SEL]}"
S_MNT="${PT_MNT[$SEL]}"

echo
echo "  Partition : $S_DEV  (${PT_FS[$SEL]}, ${PT_SIZE[$SEL]}, ${PT_FREE[$SEL]} MiB free)"
echo "  Disk      : $S_DISK   Partition number: $S_PNUM"
echo "  Mountpoint: $S_MNT"
echo
echo "  $S_DEV will be shrunk by ${SHRINK_MIB} MiB."
echo "  A new ${SHRINK_MIB} MiB BTRFS partition labeled OEMDRV will be created."
echo "  It will be mounted at $MOUNT_POINT and the repository cloned into it."
echo
read -r -p "  Type YES (uppercase) to confirm, anything else to abort: " CONFIRM
[[ "$CONFIRM" == "YES" ]] || { echo "Aborted."; exit 0; }

# ── Geometry ──────────────────────────────────────────────────────────────────

info "Reading partition geometry on $S_DISK..."

CURR_END=$(part_end_mib "$S_DISK" "$S_PNUM")
CURR_START=$(part_start_mib "$S_DISK" "$S_PNUM")
[[ -n "$CURR_END" && -n "$CURR_START" ]] || die "Could not read geometry for partition $S_PNUM on $S_DISK."

NEW_END=$(( CURR_END - SHRINK_MIB ))
NEW_FS_MIB=$(( NEW_END - CURR_START ))
OEMDRV_START=$NEW_END
OEMDRV_END=$CURR_END

echo "  Current partition : ${CURR_START} MiB – ${CURR_END} MiB"
echo "  After shrink      : ${CURR_START} MiB – ${NEW_END} MiB  (filesystem: ${NEW_FS_MIB} MiB)"
echo "  New OEMDRV        : ${OEMDRV_START} MiB – ${OEMDRV_END} MiB  (${SHRINK_MIB} MiB)"

# ── Unmount ───────────────────────────────────────────────────────────────────

WAS_MOUNTED=0
ORIG_MNT=""
if [[ "$S_MNT" != "-" && -n "$S_MNT" ]]; then
    info "Unmounting $S_DEV from $S_MNT..."
    ORIG_MNT="$S_MNT"
    umount "$S_DEV" || die "Cannot unmount $S_DEV. Close all open files and try again."
    WAS_MOUNTED=1
fi

# ── Filesystem resize ─────────────────────────────────────────────────────────

if [[ "$S_FS" == "ext4" ]]; then
    info "Running e2fsck on $S_DEV..."
    e2fsck -f "$S_DEV" || die "e2fsck found errors on $S_DEV. Fix them before retrying."

    info "Shrinking ext4 filesystem to ${NEW_FS_MIB} MiB..."
    resize2fs "$S_DEV" "${NEW_FS_MIB}M" || die "resize2fs failed."

elif [[ "$S_FS" == "btrfs" ]]; then
    info "Shrinking btrfs filesystem on $S_DEV to ${NEW_FS_MIB} MiB..."
    TMP_MNT=$(mktemp -d /tmp/oemdrv_btrfs.XXXXXX)
    mount "$S_DEV" "$TMP_MNT" || { rmdir "$TMP_MNT"; die "Cannot mount $S_DEV for btrfs resize."; }
    btrfs filesystem resize "${NEW_FS_MIB}m" "$TMP_MNT" \
        || { umount "$TMP_MNT"; rmdir "$TMP_MNT"; die "btrfs filesystem resize failed."; }
    sync
    umount "$TMP_MNT" && rmdir "$TMP_MNT"
fi

# ── Partition table update ────────────────────────────────────────────────────

info "Shrinking partition $S_PNUM to ${NEW_END} MiB in partition table..."
parted -s "$S_DISK" resizepart "$S_PNUM" "${NEW_END}MiB" \
    || die "parted resizepart failed."

info "Creating new OEMDRV partition (${OEMDRV_START}–${OEMDRV_END} MiB)..."
parted -s "$S_DISK" mkpart primary btrfs "${OEMDRV_START}MiB" "${OEMDRV_END}MiB" \
    || die "parted mkpart failed. Check that the target area is free space on $S_DISK."

partprobe "$S_DISK"
sleep 1

# Determine new partition number (highest on the disk after partprobe)
NEW_PNUM=$(parted -s "$S_DISK" -m unit MiB print 2>/dev/null \
    | awk -F: '/^[0-9]/{n=$1} END{print n}')
[[ -n "$NEW_PNUM" ]] || die "Could not determine new partition number on $S_DISK."

OEMDRV_DEV=$(new_part_device "$S_DISK" "$NEW_PNUM")
[[ -b "$OEMDRV_DEV" ]] || die "New partition device $OEMDRV_DEV not found after partprobe."

info "New partition device: $OEMDRV_DEV"

# ── Format ────────────────────────────────────────────────────────────────────

info "Formatting $OEMDRV_DEV as BTRFS (label: $OEMDRV_LABEL)..."
mkfs.btrfs -f -L "$OEMDRV_LABEL" "$OEMDRV_DEV" || die "mkfs.btrfs failed."

# ── Remount original partition ────────────────────────────────────────────────

if [[ $WAS_MOUNTED -eq 1 && -n "$ORIG_MNT" ]]; then
    info "Remounting $S_DEV to $ORIG_MNT..."
    mount "$S_DEV" "$ORIG_MNT" \
        || echo "WARNING: Could not remount $S_DEV to $ORIG_MNT — remount manually."
fi

# ── Mount OEMDRV ──────────────────────────────────────────────────────────────

info "Mounting $OEMDRV_DEV to $MOUNT_POINT (options: $MOUNT_OPTS)..."
[[ -d "$MOUNT_POINT" ]] || mkdir -p "$MOUNT_POINT"
mount -o "$MOUNT_OPTS" "$OEMDRV_DEV" "$MOUNT_POINT" || die "mount failed."

# ── Clone repository ──────────────────────────────────────────────────────────

info "Cloning $REPO_URL into $MOUNT_POINT..."
cd "$MOUNT_POINT" || die "Cannot cd to $MOUNT_POINT."
git clone --progress --depth 1 "$REPO_URL" . || die "git clone failed."

# ── Done ──────────────────────────────────────────────────────────────────────

info "Done."
echo
echo "  OEMDRV device : $OEMDRV_DEV"
echo "  Mounted at    : $MOUNT_POINT"
echo
echo "Next steps:"
echo "  1. cp $MOUNT_POINT/config/setup_system.conf.dist \\"
echo "        $MOUNT_POINT/config/setup_system.conf"
echo "  2. Edit setup_system.conf with your domain, IPA/Nextcloud FQDNs, and paths."
echo "  3. Boot the Kickstart installer — it will detect the OEMDRV partition automatically."
echo