#!/bin/bash
# Partition / filesystem setup.
# - Reformats ROOT_PART as btrfs with subvolumes
# - Mounts everything under TARGET (default /mnt)
# - Mounts existing EFI partition read-write at $TARGET/boot/efi (NEVER reformatted)

# shellcheck source=common.sh
source "$(dirname "${BASH_SOURCE[0]}")/common.sh"

setup_filesystems() {
    step "Filesystem setup on $ROOT_PART (btrfs) + EFI share $EFI_PART"

    [[ -b "$ROOT_PART" ]] || die "ROOT_PART $ROOT_PART not a block device"
    [[ -b "$EFI_PART"  ]] || die "EFI_PART  $EFI_PART not a block device"

    # Sanity: EFI must already be vfat. We never format it.
    local efi_fs
    efi_fs=$(lsblk -no FSTYPE "$EFI_PART")
    [[ "$efi_fs" == "vfat" ]] \
        || die "EFI partition $EFI_PART is '$efi_fs', expected vfat — refusing"

    # Force unmount anything currently using ROOT_PART (live ISO, prior run).
    local mp
    while read -r mp; do
        [[ -n "$mp" ]] && umount -R "$mp" 2>/dev/null || true
    done < <(lsblk -nro MOUNTPOINTS "$ROOT_PART" 2>/dev/null | grep -v '^$' || true)
    umount -R "$ROOT_PART" 2>/dev/null || true
    swapoff -a 2>/dev/null || true

    # If the live ISO's initramfs already auto-discovered a btrfs on this
    # partition (mklive runs `btrfs device scan` early), the device is
    # registered with the in-kernel btrfs module and any later mkfs sees EBUSY.
    # Forget the registration BEFORE wiping so it doesn't get re-claimed.
    btrfs device scan --forget 2>/dev/null || true

    log "wiping filesystem signatures on $ROOT_PART"
    wipefs -af "$ROOT_PART" 2>/dev/null || true
    # Zero the first 64 MiB and last 4 MiB to obliterate btrfs primary AND
    # backup superblocks. Without this, a half-written btrfs from an
    # interrupted prior run can be auto-mounted between commands.
    dd if=/dev/zero of="$ROOT_PART" bs=1M count=64 conv=fsync 2>/dev/null || true
    local part_size_b
    part_size_b=$(blockdev --getsize64 "$ROOT_PART" 2>/dev/null || echo 0)
    if [[ "$part_size_b" -gt $((8 * 1024 * 1024)) ]]; then
        local seek_mb=$(( part_size_b / 1024 / 1024 - 4 ))
        dd if=/dev/zero of="$ROOT_PART" bs=1M count=4 seek="$seek_mb" \
            conv=fsync 2>/dev/null || true
    fi
    sync
    udevadm settle 2>/dev/null || true
    btrfs device scan --forget 2>/dev/null || true

    log "lsblk view of $ROOT_PART:"
    lsblk -fno NAME,FSTYPE,LABEL,MOUNTPOINTS "$ROOT_PART" 2>&1 \
        | while read -r l; do log "  $l"; done

    log "creating btrfs on $ROOT_PART"
    # Stop udev from auto-claiming the device mid-format (race that causes
    # mkfs.btrfs to fail its final O_EXCL reopen with EBUSY even though the
    # superblock was written successfully).
    udevadm control --stop-exec-queue 2>/dev/null || true
    local mkfs_rc=0
    mkfs.btrfs -f -L void "$ROOT_PART" || mkfs_rc=$?
    sync
    udevadm control --start-exec-queue 2>/dev/null || true
    udevadm settle 2>/dev/null || true
    if [[ $mkfs_rc -ne 0 ]]; then
        # Tolerate nonzero exit if a valid btrfs was in fact written
        # (EBUSY-on-close race observed with btrfs-progs 6.11).
        local actual_fs
        actual_fs=$(blkid -o value -s TYPE "$ROOT_PART" 2>/dev/null || echo "")
        if [[ "$actual_fs" != "btrfs" ]]; then
            die "mkfs.btrfs failed on $ROOT_PART (rc=$mkfs_rc, fs='$actual_fs')"
        fi
        log "mkfs.btrfs exited rc=$mkfs_rc but blkid reports btrfs — continuing"
    fi

    local TARGET="${TARGET:-/mnt}"
    mkdir -p "$TARGET"
    mount -o "${BTRFS_MOUNT_OPTS}" "$ROOT_PART" "$TARGET"

    log "creating subvolumes: ${BTRFS_SUBVOLS[*]%%:*}"
    local entry sv
    for entry in "${BTRFS_SUBVOLS[@]}"; do
        sv="${entry%%:*}"
        # Idempotent: skip if a stale subvolume already exists from a
        # previous interrupted run (mkfs -f recreates the FS but on a fresh
        # mount the directory listing should be empty; this is defensive).
        if [[ -e "$TARGET/$sv" ]]; then
            log "  subvolume $sv already present — skipping"
            continue
        fi
        btrfs subvolume create "$TARGET/$sv"
    done
    umount "$TARGET"

    log "remounting subvolumes"
    # First mount the root subvolume (must be the first entry, conventionally @).
    local root_entry="${BTRFS_SUBVOLS[0]}"
    local root_sv="${root_entry%%:*}"
    mount -o "${BTRFS_MOUNT_OPTS},subvol=${root_sv}" "$ROOT_PART" "$TARGET"

    # Pre-create mountpoints + EFI dir.
    mkdir -p "$TARGET/boot/efi"
    for entry in "${BTRFS_SUBVOLS[@]:1}"; do
        local mp="${entry##*:}"
        mkdir -p "$TARGET$mp"
    done

    for entry in "${BTRFS_SUBVOLS[@]:1}"; do
        sv="${entry%%:*}"
        local mp="${entry##*:}"
        mount -o "${BTRFS_MOUNT_OPTS},subvol=${sv}" "$ROOT_PART" "$TARGET$mp"
    done

    # Mount shared EFI WITHOUT formatting.
    mount "$EFI_PART" "$TARGET/boot/efi"

    # Make sure we're not about to clobber Windows' EFI loader.
    if [[ -d "$TARGET/boot/efi/EFI/Microsoft" ]]; then
        log "Windows EFI files detected — will preserve EFI/Microsoft/* untouched"
    else
        warn "no EFI/Microsoft dir found on $EFI_PART — proceeding anyway"
    fi

    ok "filesystems mounted under $TARGET"
    findmnt -R "$TARGET" >> "$LOG_FILE"

    export TARGET
}