#!/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.

[[ "$EUID" -eq 0 ]] || { echo "ERROR: This script must be run as root." >&2; exit 1; }

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

# Parallel arrays for unpartitioned free space regions, indexed 0..FS_IDX-1
FS_IDX=0
FS_DISK=()   # /dev/sda
FS_START=()  # region start, MiB
FS_END=()    # region end, MiB
FS_SIZE=()   # region size, MiB

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)
}

collect_free_space() {
    local disk
    while IFS= read -r disk; do
        [[ -b "$disk" ]] || continue
        while read -r s e sz; do
            FS_DISK[$FS_IDX]="$disk"
            FS_START[$FS_IDX]="$s"
            FS_END[$FS_IDX]="$e"
            FS_SIZE[$FS_IDX]="$sz"
            (( FS_IDX++ )) || true
        done < <(
            parted -s "$disk" -m unit MiB print free 2>/dev/null \
            | awk -F'[;:]' -v min="$SHRINK_MIB" '
                $1+0 > 0 {
                    for (i = 1; i <= NF; i++) {
                        if ($i == "free") {
                            start=$2; end=$3; size=$4;
                            gsub(/MiB/,"",start); gsub(/MiB/,"",end); gsub(/MiB/,"",size);
                            s=int(start+0); e=int(end+0); sz=int(size+0);
                            if (sz >= min) print s " " e " " sz;
                            break
                        }
                    }
                }'
        )
    done < <(lsblk -d -p -n -o NAME 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
}

show_free_table() {
    [[ $FS_IDX -gt 0 ]] || return
    echo
    echo "Unpartitioned free space regions (>= ${SHRINK_MIB} MiB):"
    hr
    printf "  %-4s  %-14s  %-12s  %-12s  %s\n" "#" "Disk" "Start (MiB)" "End (MiB)" "Size (MiB)"
    hr
    for (( i=0; i<FS_IDX; i++ )); do
        printf "  %-4s  %-14s  %-12s  %-12s  %s\n" \
            "$((i+1))" \
            "${FS_DISK[$i]}" \
            "${FS_START[$i]}" \
            "${FS_END[$i]}" \
            "${FS_SIZE[$i]}"
    done
    hr
    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 for shrinkable partitions and unpartitioned free space..."
collect_partitions
collect_free_space

[[ $PT_IDX -gt 0 || $FS_IDX -gt 0 ]] \
    || die "No ext4 or btrfs partitions and no free disk space found on this system."

# Count shrinkable partitions
shrink_count=0
for (( i=0; i<PT_IDX; i++ )); do [[ "${PT_OK[$i]}" == "1" ]] && (( shrink_count++ )) || true; done

[[ $PT_IDX -gt 0 ]] && show_table
show_free_table

[[ $shrink_count -gt 0 || $FS_IDX -gt 0 ]] \
    || die "No shrinkable partitions and no free space found. A non-root ext4 or btrfs partition needs at least ${MIN_FREE_MIB} MiB free."

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

MODE=""   # "freespace" | "shrink"
SEL=-1

while true; do
    echo
    if [[ $FS_IDX -gt 0 && $shrink_count -gt 0 ]]; then
        read -r -p "Enter  f<n>  to use free space,  s<n>  to shrink a partition, or  q  to quit: " INPUT
    elif [[ $FS_IDX -gt 0 ]]; then
        read -r -p "Enter number of free space region to use, or  q  to quit: " INPUT
        [[ "$INPUT" =~ ^[0-9]+$ ]] && INPUT="f${INPUT}"
    else
        read -r -p "Enter number of partition to shrink, or  q  to quit: " INPUT
        [[ "$INPUT" =~ ^[0-9]+$ ]] && INPUT="s${INPUT}"
    fi

    [[ "$INPUT" =~ ^[qQ]$ ]] && { echo "Aborted."; exit 0; }

    if [[ "$INPUT" =~ ^[fF]([0-9]+)$ ]]; then
        local_i=$(( ${BASH_REMATCH[1]} - 1 ))
        (( local_i >= 0 && local_i < FS_IDX )) || { echo "  Number out of range."; continue; }
        MODE="freespace"; SEL=$local_i; break

    elif [[ "$INPUT" =~ ^[sS]([0-9]+)$ ]]; then
        local_i=$(( ${BASH_REMATCH[1]} - 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; }
        MODE="shrink"; SEL=$local_i; break

    else
        echo "  Invalid input."
    fi
done

# ── Confirm ───────────────────────────────────────────────────────────────────

echo
if [[ "$MODE" == "freespace" ]]; then
    WORK_DISK="${FS_DISK[$SEL]}"
    OEMDRV_START="${FS_START[$SEL]}"
    OEMDRV_END=$(( FS_START[$SEL] + SHRINK_MIB ))
    echo "  Disk        : $WORK_DISK"
    echo "  Free region : ${FS_START[$SEL]} MiB – ${FS_END[$SEL]} MiB  (${FS_SIZE[$SEL]} MiB available)"
    echo "  New OEMDRV  : ${OEMDRV_START} MiB – ${OEMDRV_END} MiB  (${SHRINK_MIB} MiB)"
    echo "  It will be mounted at $MOUNT_POINT and the repository cloned into it."
else
    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]}"
    WORK_DISK="$S_DISK"
    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."
fi
echo
read -r -p "  Type YES (uppercase) to confirm, anything else to abort: " CONFIRM
[[ "$CONFIRM" == "YES" ]] || { echo "Aborted."; exit 0; }

# ── Shrink path: geometry + filesystem resize + partition table update ─────────

WAS_MOUNTED=0
ORIG_MNT=""

if [[ "$MODE" == "shrink" ]]; then
    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)"

    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

    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

    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."
fi

# ── Create OEMDRV partition ───────────────────────────────────────────────────

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

partprobe "$WORK_DISK"
sleep 1

# Determine new partition number (highest on the disk after partprobe)
NEW_PNUM=$(parted -s "$WORK_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 $WORK_DISK."

OEMDRV_DEV=$(new_part_device "$WORK_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 (shrink path only) ─────────────────────────────

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