#!/bin/bash

# init.sh for Legion S7 15ACH6
# Gentoo bootstrap script with Btrfs subvolumes

# --- Script Configuration Variables ---
GENTOO_MIRROR_BASE="https://mirrors.rit.edu/gentoo/"
GENTOO_FALLBACK_MIRRORS=(
    "https://distfiles.gentoo.org/"
    "https://gentoo.osuosl.org/"
    "https://mirror.bytemark.co.uk/gentoo/"
)

# Default swap size in GB (reduced from 32GB to 16GB)
SWAP_SIZE_GB=16

# Init system: "openrc" or "systemd"
INIT_SYSTEM="openrc"

# --- Colors for Output ---
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
NC='\033[0m'

set -euo pipefail

# --- Cleanup Trap ---
cleanup() {
    local exit_code=$?
    if [[ $exit_code -ne 0 ]] && [[ -n "${GENTOO_ROOT_MOUNT_GLOBAL:-}" ]] && mountpoint -q "$GENTOO_ROOT_MOUNT_GLOBAL" 2>/dev/null; then
        echo -e "${YELLOW}Script interrupted. Cleaning up mounts...${NC}"
        umount -R "$GENTOO_ROOT_MOUNT_GLOBAL" 2>/dev/null || true
    fi
    # Close LUKS container if open
    if [[ -e "$CRYPT_ROOT_GLOBAL" ]]; then
        echo -e "${YELLOW}Closing LUKS container...${NC}"
        cryptsetup close cryptroot 2>/dev/null || true
    fi
    exit $exit_code
}
trap cleanup EXIT

# --- Global Variables ---
TARGET_DISK_GLOBAL=""
DISK_SIZE_GB_GLOBAL=""
GENTOO_ROOT_MOUNT_GLOBAL="/mnt/gentoo"
BTRFS_TEMP_MOUNT_GLOBAL="/mnt/btrfs_temp"
PART1_GLOBAL=""
PART2_GLOBAL=""
PART3_GLOBAL=""
CRYPT_ROOT_GLOBAL="/dev/mapper/cryptroot"

# Btrfs mount options (removed redundant ssd and space_cache=v2)
BTRFS_OPTS="noatime,compress=zstd"

# --- Helper Functions ---

get_disk_size() {
    local disk="$1"
    local size_bytes
    size_bytes=$(lsblk -b -n -o SIZE "$disk" | head -n 1)
    if [ -z "$size_bytes" ]; then
        echo "unknown"
        return 1
    fi
    awk -v bytes="$size_bytes" 'BEGIN { printf "%.2f", bytes / (1024^3) }'
    return 0
}

force_unmount_and_swapoff_target_disk() {
    echo -e "${YELLOW}Attempting to unmount all partitions and deactivate swap on ${TARGET_DISK_GLOBAL}...${NC}"

    local partitions=()
    mapfile -t partitions < <(lsblk -rno NAME "$TARGET_DISK_GLOBAL" | tail -n +2)

    local partition_devices=()
    for part in "${partitions[@]}"; do
        partition_devices+=("/dev/$part")
    done

    # Deactivate swap
    echo -e "${YELLOW}Checking for active swap partitions...${NC}"
    for device in "${partition_devices[@]}"; do
        if swapon --show=NAME --noheadings | grep -q "^$device$"; then
            echo -e "${YELLOW}  Deactivating swap on ${device}...${NC}"
            if swapoff "$device" 2>/dev/null; then
                echo -e "${GREEN}    Successfully deactivated swap on ${device}${NC}"
            else
                echo -e "${RED}    Warning: Could not deactivate swap on ${device}${NC}"
            fi
        fi
    done

    # Unmount partitions (multiple attempts)
    echo -e "${YELLOW}Checking for mounted partitions...${NC}"
    for attempt in {1..3}; do
        local mounted_found=false

        for device in "${partition_devices[@]}"; do
            if mount | grep -q "^$device "; then
                mounted_found=true
                local mountpoint
                mountpoint=$(mount | grep "^$device " | awk '{print $3}' | head -n1)
                echo -e "${YELLOW}  Attempt ${attempt}: Unmounting ${device} from ${mountpoint}...${NC}"

                if [ "$attempt" -eq 1 ]; then
                    umount "$device" 2>/dev/null && echo -e "${GREEN}    Successfully unmounted ${device}${NC}" && continue
                elif [ "$attempt" -eq 2 ]; then
                    umount -f "$device" 2>/dev/null && echo -e "${GREEN}    Force unmounted ${device}${NC}" && continue
                else
                    umount -l "$device" 2>/dev/null && echo -e "${YELLOW}    Lazy unmounted ${device}${NC}" && continue
                fi
                echo -e "${RED}    Warning: Could not unmount ${device}${NC}"
            fi
        done

        if ! $mounted_found; then
            break
        fi
        sleep 1
    done

    partprobe "$TARGET_DISK_GLOBAL" 2>/dev/null || true
    sync

    echo -e "${GREEN}Verification after unmount/swapoff operations:${NC}"
    lsblk "$TARGET_DISK_GLOBAL" || true
    mount | grep "$TARGET_DISK_GLOBAL" || echo "  No active mounts found for ${TARGET_DISK_GLOBAL}"
    swapon --show | grep "$TARGET_DISK_GLOBAL" || echo "  No active swap found for ${TARGET_DISK_GLOBAL}"

    return 0
}

wipe_disk_signatures() {
    echo -e "${YELLOW}Wiping existing filesystem signatures and partition table...${NC}"

    wipefs -af "$TARGET_DISK_GLOBAL" 2>/dev/null || true

    echo -e "${YELLOW}Zeroing beginning and end of disk...${NC}"
    dd if=/dev/zero of="$TARGET_DISK_GLOBAL" bs=1M count=10 2>/dev/null || true

    local disk_size_sectors
    disk_size_sectors=$(blockdev --getsz "$TARGET_DISK_GLOBAL" 2>/dev/null || echo "0")
    if [ "$disk_size_sectors" -gt 20480 ]; then
        local last_mb_start=$((disk_size_sectors - 20480))
        dd if=/dev/zero of="$TARGET_DISK_GLOBAL" bs=512 seek="$last_mb_start" count=20480 2>/dev/null || true
    fi

    partprobe "$TARGET_DISK_GLOBAL" 2>/dev/null || true
    sync
    echo -e "${GREEN}Disk wiping complete${NC}"
}

# --- Main Functions ---

select_target_disk() {
    local disks=()
    local i=1
    echo -e "${YELLOW}Identifying available disks...${NC}"
    lsblk -p -o NAME,SIZE,MODEL,VENDOR,TYPE

    local root_dev
    root_dev=$(findmnt -n -o SOURCE / | awk '{print $1}')
    local boot_disk_basename
    if [[ "$root_dev" == /dev/sd* ]]; then
        boot_disk_basename=$(basename "${root_dev%[0-9]*}")
    elif [[ "$root_dev" == /dev/nvme* ]]; then
        boot_disk_basename=$(basename "${root_dev%p[0-9]*}")
    else
        boot_disk_basename=$(basename "$root_dev")
    fi

    local all_disks_names=()
    mapfile -t all_disks_names < <(lsblk -dno NAME | sort)

    local non_boot_disks_filtered=()
    for d_name in "${all_disks_names[@]}"; do
        if [[ "$d_name" != "$boot_disk_basename" ]]; then
            if [[ -b "/dev/${d_name}" ]]; then
                non_boot_disks_filtered+=("$d_name")
            fi
        fi
    done

    if [ ${#non_boot_disks_filtered[@]} -eq 0 ]; then
        echo -e "${RED}No suitable internal disks found for installation. Aborting.${NC}"
        return 1
    fi

    echo -e "${YELLOW}Select your TARGET DISK:${NC}"
    for disk_name in "${non_boot_disks_filtered[@]}"; do
        local current_disk_path="/dev/${disk_name}"
        local size_val
        size_val=$(get_disk_size "$current_disk_path")
        local model_val
        model_val=$(lsblk -dno MODEL "$current_disk_path" | head -n 1)

        disks+=("$current_disk_path")
        echo "  ${i}) ${current_disk_path} - ${size_val} GB (${model_val})"
        ((i++))
    done

    local choice
    while true; do
        read -r -p "Enter the number corresponding to your target disk: " choice
        if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le ${#disks[@]} ]; then
            TARGET_DISK_GLOBAL="${disks[choice-1]}"
            DISK_SIZE_GB_GLOBAL=$(get_disk_size "$TARGET_DISK_GLOBAL")
            echo -e "${YELLOW}Selected: ${TARGET_DISK_GLOBAL} (${DISK_SIZE_GB_GLOBAL} GB)${NC}"
            read -r -p "Are you ABSOLUTELY SURE? This will ERASE ALL DATA! (type 'yes' to confirm): " confirmation
            if [[ "$confirmation" == "yes" ]]; then
                echo -e "${GREEN}Target disk confirmed.${NC}"
                return 0
            else
                echo -e "${RED}Disk selection cancelled.${NC}"
                return 1
            fi
        else
            echo -e "${RED}Invalid selection.${NC}"
        fi
    done
}

partition_and_format_disk() {
    echo -e "${YELLOW}--- Disk Partitioning & Formatting ---${NC}"
    echo -e "Target disk: ${TARGET_DISK_GLOBAL} (${DISK_SIZE_GB_GLOBAL} GB)"
    echo -e "Swap size: ${SWAP_SIZE_GB} GB"

    # Check for UEFI
    if [ ! -d "/sys/firmware/efi/efivars" ]; then
        echo -e "${RED}ERROR: UEFI boot mode not detected. This script requires UEFI.${NC}"
        return 1
    fi
    echo -e "${GREEN}Detected UEFI boot mode.${NC}"

    if ! force_unmount_and_swapoff_target_disk; then
        echo -e "${RED}Error during pre-partition cleanup. Aborting.${NC}"
        return 1
    fi

    wipe_disk_signatures

    # Create GPT table with 1MiB alignment (optimal for SSDs/NVMe)
    if ! sgdisk --clear --mbrtogpt --set-alignment=2048 "$TARGET_DISK_GLOBAL"; then
        echo -e "${RED}Error creating GPT table. Aborting.${NC}"
        return 1
    fi

    partprobe "$TARGET_DISK_GLOBAL" 2>/dev/null || true
    sleep 2

    # Partition 1: EFI (512MB)
    if ! sgdisk --set-alignment=2048 -n 1:0:+512MiB -t 1:EF00 -c 1:"EFI System Partition" "$TARGET_DISK_GLOBAL"; then
        echo -e "${RED}Error creating EFI partition. Aborting.${NC}"
        return 1
    fi

    # Partition 2: Swap
    if ! sgdisk --set-alignment=2048 -n 2:0:+${SWAP_SIZE_GB}GiB -t 2:8200 -c 2:"Linux Swap" "$TARGET_DISK_GLOBAL"; then
        echo -e "${RED}Error creating swap partition. Aborting.${NC}"
        return 1
    fi

    # Partition 3: Btrfs Root
    if ! sgdisk --set-alignment=2048 -n 3:0:0 -t 3:8300 -c 3:"Linux Btrfs Root" "$TARGET_DISK_GLOBAL"; then
        echo -e "${RED}Error creating Btrfs partition. Aborting.${NC}"
        return 1
    fi

    echo -e "${GREEN}Disk partitioning complete.${NC}"
    sgdisk -p "$TARGET_DISK_GLOBAL"

    # Wait for kernel to recognize partitions
    partprobe "$TARGET_DISK_GLOBAL" 2>/dev/null || true
    sleep 3

    # Determine partition naming scheme
    if [[ "$TARGET_DISK_GLOBAL" == /dev/nvme* ]]; then
        PART1_GLOBAL="${TARGET_DISK_GLOBAL}p1"
        PART2_GLOBAL="${TARGET_DISK_GLOBAL}p2"
        PART3_GLOBAL="${TARGET_DISK_GLOBAL}p3"
    elif [[ "$TARGET_DISK_GLOBAL" == /dev/mmcblk* ]]; then
        PART1_GLOBAL="${TARGET_DISK_GLOBAL}p1"
        PART2_GLOBAL="${TARGET_DISK_GLOBAL}p2"
        PART3_GLOBAL="${TARGET_DISK_GLOBAL}p3"
    else
        PART1_GLOBAL="${TARGET_DISK_GLOBAL}1"
        PART2_GLOBAL="${TARGET_DISK_GLOBAL}2"
        PART3_GLOBAL="${TARGET_DISK_GLOBAL}3"
    fi

    echo -e "${YELLOW}Detected partitions: ${PART1_GLOBAL}, ${PART2_GLOBAL}, ${PART3_GLOBAL}${NC}"

    # Wait for partitions
    for part in "$PART1_GLOBAL" "$PART2_GLOBAL" "$PART3_GLOBAL"; do
        local wait_count=0
        while [ ! -b "$part" ] && [ $wait_count -lt 10 ]; do
            echo -e "${YELLOW}Waiting for ${part}...${NC}"
            sleep 1
            partprobe "$TARGET_DISK_GLOBAL" 2>/dev/null || true
            ((wait_count++))
        done
        if [ ! -b "$part" ]; then
            echo -e "${RED}Error: Partition ${part} not found. Aborting.${NC}"
            return 1
        fi
    done

    echo -e "${YELLOW}Formatting partitions...${NC}"

    # Format EFI
    if ! mkfs.fat -F 32 "$PART1_GLOBAL"; then
        echo -e "${RED}Error formatting EFI partition. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}Formatted ${PART1_GLOBAL} as FAT32.${NC}"

    # Skip swap formatting - will be encrypted with random key at boot via fstab
    echo -e "${YELLOW}Swap will be encrypted at boot (skipping format now)${NC}"

    # Setup LUKS2 encryption on root partition
    echo -e "${YELLOW}Setting up LUKS2 encryption on ${PART3_GLOBAL}...${NC}"
    echo -e "${GREEN}>>> Use OnlyKey to enter your LUKS passphrase <<<${NC}"
    if ! cryptsetup luksFormat --type luks2 --cipher aes-xts-plain64 --key-size 512 --hash sha512 "$PART3_GLOBAL"; then
        echo -e "${RED}Error creating LUKS container. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}LUKS container created.${NC}"

    echo -e "${YELLOW}Opening LUKS container...${NC}"
    echo -e "${GREEN}>>> Enter passphrase again to open LUKS <<<${NC}"
    if ! cryptsetup open "$PART3_GLOBAL" cryptroot; then
        echo -e "${RED}Error opening LUKS container. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}LUKS container opened at ${CRYPT_ROOT_GLOBAL}${NC}"

    # Format Btrfs on LUKS container
    if ! mkfs.btrfs -f "$CRYPT_ROOT_GLOBAL"; then
        echo -e "${RED}Error formatting Btrfs on LUKS. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}Formatted ${CRYPT_ROOT_GLOBAL} as Btrfs.${NC}"

    echo -e "${YELLOW}Creating Btrfs subvolumes...${NC}"

    mkdir -p "$BTRFS_TEMP_MOUNT_GLOBAL"
    if ! mount "$CRYPT_ROOT_GLOBAL" "$BTRFS_TEMP_MOUNT_GLOBAL"; then
        echo -e "${RED}Error mounting Btrfs for subvolume creation. Aborting.${NC}"
        return 1
    fi

    # Create subvolumes
    local subvolumes=("@" "@home" "@var" "@log" "@snapshots")
    for subvol in "${subvolumes[@]}"; do
        if ! btrfs subvolume create "${BTRFS_TEMP_MOUNT_GLOBAL}/${subvol}"; then
            echo -e "${RED}Error creating ${subvol} subvolume. Aborting.${NC}"
            umount "$BTRFS_TEMP_MOUNT_GLOBAL" 2>/dev/null || true
            return 1
        fi
        echo -e "${GREEN}Created subvolume: ${subvol}${NC}"
    done

    if ! umount "$BTRFS_TEMP_MOUNT_GLOBAL"; then
        echo -e "${RED}Error unmounting Btrfs temp. Aborting.${NC}"
        return 1
    fi
    rmdir "$BTRFS_TEMP_MOUNT_GLOBAL" 2>/dev/null || true
    echo -e "${GREEN}Btrfs subvolumes created.${NC}"

    return 0
}

mount_partitions() {
    echo -e "${YELLOW}--- Mounting Partitions ---${NC}"

    if mountpoint -q "$GENTOO_ROOT_MOUNT_GLOBAL" 2>/dev/null; then
        echo -e "${YELLOW}${GENTOO_ROOT_MOUNT_GLOBAL} already mounted, unmounting...${NC}"
        umount -R "$GENTOO_ROOT_MOUNT_GLOBAL" 2>/dev/null || true
    fi

    mkdir -p "$GENTOO_ROOT_MOUNT_GLOBAL"

    # Mount @ subvolume (root)
    if ! mount -o "${BTRFS_OPTS},subvol=@" "$CRYPT_ROOT_GLOBAL" "$GENTOO_ROOT_MOUNT_GLOBAL"; then
        echo -e "${RED}Error mounting @ subvolume. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}Mounted @ to ${GENTOO_ROOT_MOUNT_GLOBAL}${NC}"

    # Mount @home
    mkdir -p "${GENTOO_ROOT_MOUNT_GLOBAL}/home"
    if ! mount -o "${BTRFS_OPTS},subvol=@home" "$CRYPT_ROOT_GLOBAL" "${GENTOO_ROOT_MOUNT_GLOBAL}/home"; then
        echo -e "${RED}Error mounting @home subvolume. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}Mounted @home to ${GENTOO_ROOT_MOUNT_GLOBAL}/home${NC}"

    # Mount @var
    mkdir -p "${GENTOO_ROOT_MOUNT_GLOBAL}/var"
    if ! mount -o "${BTRFS_OPTS},subvol=@var" "$CRYPT_ROOT_GLOBAL" "${GENTOO_ROOT_MOUNT_GLOBAL}/var"; then
        echo -e "${RED}Error mounting @var subvolume. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}Mounted @var to ${GENTOO_ROOT_MOUNT_GLOBAL}/var${NC}"

    # Mount @log
    mkdir -p "${GENTOO_ROOT_MOUNT_GLOBAL}/var/log"
    if ! mount -o "${BTRFS_OPTS},subvol=@log" "$CRYPT_ROOT_GLOBAL" "${GENTOO_ROOT_MOUNT_GLOBAL}/var/log"; then
        echo -e "${RED}Error mounting @log subvolume. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}Mounted @log to ${GENTOO_ROOT_MOUNT_GLOBAL}/var/log${NC}"

    # Mount @snapshots
    mkdir -p "${GENTOO_ROOT_MOUNT_GLOBAL}/.snapshots"
    if ! mount -o "${BTRFS_OPTS},subvol=@snapshots" "$CRYPT_ROOT_GLOBAL" "${GENTOO_ROOT_MOUNT_GLOBAL}/.snapshots"; then
        echo -e "${RED}Error mounting @snapshots subvolume. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}Mounted @snapshots to ${GENTOO_ROOT_MOUNT_GLOBAL}/.snapshots${NC}"

    # Mount EFI
    mkdir -p "${GENTOO_ROOT_MOUNT_GLOBAL}/boot"
    if ! mount "$PART1_GLOBAL" "${GENTOO_ROOT_MOUNT_GLOBAL}/boot"; then
        echo -e "${RED}Error mounting EFI partition. Aborting.${NC}"
        return 1
    fi
    echo -e "${GREEN}Mounted EFI to ${GENTOO_ROOT_MOUNT_GLOBAL}/boot${NC}"

    echo -e "${GREEN}All partitions mounted successfully!${NC}"
    echo -e "${YELLOW}Mount verification:${NC}"
    findmnt -R "$GENTOO_ROOT_MOUNT_GLOBAL"
    return 0
}

download_and_extract_stage3() {
    echo -e "${YELLOW}--- Downloading & Extracting Stage 3 ---${NC}"

    local stage3_filename=""
    local digests_filename=""
    local stage3_base_url=""
    local successful_mirror=""

    # Determine stage3 variant based on init system
    local stage3_variant="stage3-amd64-${INIT_SYSTEM}"
    echo -e "${YELLOW}Init system: ${INIT_SYSTEM} (variant: ${stage3_variant})${NC}"

    local all_mirrors=("$GENTOO_MIRROR_BASE" "${GENTOO_FALLBACK_MIRRORS[@]}")

    for mirror in "${all_mirrors[@]}"; do
        stage3_base_url="${mirror}releases/amd64/autobuilds/current-${stage3_variant}/"
        echo -e "${YELLOW}Trying mirror: ${stage3_base_url}${NC}"

        if ! curl -s --connect-timeout 10 "${stage3_base_url}" >/dev/null; then
            echo -e "${RED}Cannot connect to ${mirror}, trying next...${NC}"
            continue
        fi

        echo -e "${YELLOW}Fetching current Stage 3 filename...${NC}"
        local latest_file="${stage3_base_url}latest-${stage3_variant}.txt"
        local latest_content

        if ! latest_content=$(curl -s --fail "$latest_file"); then
            echo -e "${RED}Could not retrieve latest file from ${mirror}, trying next...${NC}"
            continue
        fi

        # Parse PGP-signed content
        stage3_filename=$(echo "$latest_content" | \
            sed -n '/-----BEGIN PGP SIGNED MESSAGE-----/,/-----BEGIN PGP SIGNATURE-----/p' | \
            grep '\.tar\.xz' | \
            awk '{print $1}' | \
            head -n 1)

        # Fallback: direct pattern matching
        if [ -z "$stage3_filename" ]; then
            stage3_filename=$(echo "$latest_content" | grep -oE "${stage3_variant}-[0-9T]+Z?\.tar\.xz" | head -n 1)
        fi

        if [ -z "$stage3_filename" ] || [[ "$stage3_filename" != *.tar.xz ]]; then
            echo -e "${RED}Could not parse stage3 filename from ${mirror}, trying next...${NC}"
            continue
        fi

        digests_filename="${stage3_filename}.DIGESTS"
        successful_mirror="$mirror"
        echo -e "${GREEN}Found Stage 3: ${stage3_filename}${NC}"
        break
    done

    if [ -z "$successful_mirror" ]; then
        echo -e "${RED}Error: Could not find working mirror. Check network connection.${NC}"
        return 1
    fi

    stage3_base_url="${successful_mirror}releases/amd64/autobuilds/current-${stage3_variant}/"

    # Download in target directory
    pushd "$GENTOO_ROOT_MOUNT_GLOBAL" > /dev/null || {
        echo -e "${RED}Error: Could not change to ${GENTOO_ROOT_MOUNT_GLOBAL}. Aborting.${NC}"
        return 1
    }

    # Verify files exist
    echo -e "${YELLOW}Verifying files on mirror...${NC}"
    if ! curl -s --head "${stage3_base_url}${stage3_filename}" | grep -q "200 OK"; then
        echo -e "${RED}Error: Stage 3 file not found on mirror.${NC}"
        popd > /dev/null
        return 1
    fi

    # Download stage3
    echo -e "${YELLOW}Downloading ${stage3_filename}...${NC}"
    if ! wget --progress=bar:force --timeout=30 --tries=3 -c "${stage3_base_url}${stage3_filename}"; then
        echo -e "${RED}Error downloading Stage 3 tarball. Aborting.${NC}"
        popd > /dev/null
        return 1
    fi

    # Download DIGESTS
    echo -e "${YELLOW}Downloading ${digests_filename}...${NC}"
    if ! wget --progress=bar:force --timeout=30 --tries=3 -c "${stage3_base_url}${digests_filename}"; then
        echo -e "${RED}Error downloading DIGESTS file. Aborting.${NC}"
        popd > /dev/null
        return 1
    fi

    # Verify checksum
    echo -e "${YELLOW}Verifying checksum...${NC}"
    local checksum_output
    checksum_output=$(sha512sum -c --ignore-missing "$digests_filename" 2>&1)
    echo "$checksum_output"

    if echo "$checksum_output" | grep -q "${stage3_filename}: OK"; then
        echo -e "${GREEN}Stage 3 checksum VERIFIED.${NC}"
    else
        echo -e "${RED}Stage 3 checksum verification FAILED.${NC}"
        popd > /dev/null
        return 1
    fi

    # Check disk space before extraction (need at least 2GB)
    echo -e "${YELLOW}Checking available disk space...${NC}"
    local available_kb
    available_kb=$(df -k --output=avail "$GENTOO_ROOT_MOUNT_GLOBAL" | tail -1)
    if [ "$available_kb" -lt 2000000 ]; then
        echo -e "${RED}Error: Not enough disk space for extraction (need 2GB+, have ${available_kb}KB).${NC}"
        popd > /dev/null
        return 1
    fi
    echo -e "${GREEN}Disk space OK (${available_kb}KB available)${NC}"

    # Extract
    echo -e "${YELLOW}Extracting Stage 3 tarball...${NC}"
    if ! tar xpf "$stage3_filename" --xattrs-include='*.*' --numeric-owner; then
        echo -e "${RED}Error extracting Stage 3 tarball. Aborting.${NC}"
        popd > /dev/null
        return 1
    fi
    echo -e "${GREEN}Stage 3 extraction complete.${NC}"

    rm -f "$stage3_filename" "$digests_filename"
    echo -e "${GREEN}Cleaned up downloaded files.${NC}"

    popd > /dev/null
    return 0
}

prepare_chroot() {
    echo -e "${YELLOW}--- Preparing Chroot Environment ---${NC}"

    # Copy DNS info
    cp --dereference /etc/resolv.conf "${GENTOO_ROOT_MOUNT_GLOBAL}/etc/"
    echo -e "${GREEN}Copied resolv.conf${NC}"

    # Mount necessary filesystems
    mount --rbind /proc "${GENTOO_ROOT_MOUNT_GLOBAL}/proc"
    mount --make-rslave "${GENTOO_ROOT_MOUNT_GLOBAL}/proc"
    echo -e "${GREEN}Mounted /proc${NC}"

    mount --rbind /sys "${GENTOO_ROOT_MOUNT_GLOBAL}/sys"
    mount --make-rslave "${GENTOO_ROOT_MOUNT_GLOBAL}/sys"
    echo -e "${GREEN}Mounted /sys${NC}"

    mount --rbind /dev "${GENTOO_ROOT_MOUNT_GLOBAL}/dev"
    mount --make-rslave "${GENTOO_ROOT_MOUNT_GLOBAL}/dev"
    echo -e "${GREEN}Mounted /dev${NC}"

    mount --rbind /run "${GENTOO_ROOT_MOUNT_GLOBAL}/run"
    mount --make-rslave "${GENTOO_ROOT_MOUNT_GLOBAL}/run"
    echo -e "${GREEN}Mounted /run${NC}"

    echo -e "${GREEN}Chroot environment prepared.${NC}"
    return 0
}

# --- Portage Configuration Copy ---

copy_portage_config() {
    echo -e "${YELLOW}--- Copying Portage Configuration ---${NC}"

    local script_dir
    script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    local portage_dir="${GENTOO_ROOT_MOUNT_GLOBAL}/etc/portage"

    # Copy make.conf
    if [[ -f "${script_dir}/make.conf" ]]; then
        cp "${script_dir}/make.conf" "${portage_dir}/make.conf"
        echo -e "${GREEN}Copied make.conf${NC}"
    else
        echo -e "${RED}Warning: make.conf not found in ${script_dir}${NC}"
    fi

    # Copy ccache.conf
    if [[ -f "${script_dir}/ccache.conf" ]]; then
        cp "${script_dir}/ccache.conf" "${GENTOO_ROOT_MOUNT_GLOBAL}/etc/ccache.conf"
        echo -e "${GREEN}Copied ccache.conf${NC}"
    fi

    # Copy package.accept_keywords/
    if [[ -d "${script_dir}/package.accept_keywords" ]]; then
        mkdir -p "${portage_dir}/package.accept_keywords"
        cp -r "${script_dir}/package.accept_keywords/"* "${portage_dir}/package.accept_keywords/" 2>/dev/null || true
        echo -e "${GREEN}Copied package.accept_keywords/${NC}"
    fi

    # Copy package.use/
    if [[ -d "${script_dir}/package.use" ]]; then
        mkdir -p "${portage_dir}/package.use"
        cp -r "${script_dir}/package.use/"* "${portage_dir}/package.use/" 2>/dev/null || true
        echo -e "${GREEN}Copied package.use/${NC}"
    fi

    # Copy sets/
    if [[ -d "${script_dir}/sets" ]]; then
        mkdir -p "${portage_dir}/sets"
        cp -r "${script_dir}/sets/"* "${portage_dir}/sets/" 2>/dev/null || true
        echo -e "${GREEN}Copied sets/${NC}"
    fi

    # Copy dracut.conf.d/
    if [[ -d "${script_dir}/dracut.conf.d" ]]; then
        mkdir -p "${GENTOO_ROOT_MOUNT_GLOBAL}/etc/dracut.conf.d"
        cp -r "${script_dir}/dracut.conf.d/"* "${GENTOO_ROOT_MOUNT_GLOBAL}/etc/dracut.conf.d/" 2>/dev/null || true
        echo -e "${GREEN}Copied dracut.conf.d/${NC}"
    fi

    return 0
}

# --- LUKS & Filesystem Configuration ---

generate_crypttab() {
    echo -e "${YELLOW}Generating /etc/crypttab...${NC}"
    local luks_uuid
    luks_uuid=$(blkid -s UUID -o value "$PART3_GLOBAL")
    echo "cryptroot UUID=${luks_uuid} none luks" > "${GENTOO_ROOT_MOUNT_GLOBAL}/etc/crypttab"
    echo -e "${GREEN}Generated /etc/crypttab${NC}"
}

generate_fstab() {
    echo -e "${YELLOW}Generating /etc/fstab...${NC}"
    local efi_uuid
    efi_uuid=$(blkid -s UUID -o value "$PART1_GLOBAL")

    cat > "${GENTOO_ROOT_MOUNT_GLOBAL}/etc/fstab" << EOF
# /etc/fstab - Generated by init.sh
# Gentoo v3 Legion Laptop with LUKS encryption

# EFI System Partition
UUID=${efi_uuid}  /boot  vfat  noatime,defaults  0 2

# Encrypted swap (random key each boot - no hibernate support)
${PART2_GLOBAL}  none  swap  sw,cipher=aes-xts-plain64,size=256  0 0

# Btrfs on LUKS (cryptroot)
${CRYPT_ROOT_GLOBAL}  /           btrfs  ${BTRFS_OPTS},subvol=@           0 0
${CRYPT_ROOT_GLOBAL}  /home       btrfs  ${BTRFS_OPTS},subvol=@home       0 0
${CRYPT_ROOT_GLOBAL}  /var        btrfs  ${BTRFS_OPTS},subvol=@var        0 0
${CRYPT_ROOT_GLOBAL}  /var/log    btrfs  ${BTRFS_OPTS},subvol=@log        0 0
${CRYPT_ROOT_GLOBAL}  /.snapshots btrfs  ${BTRFS_OPTS},subvol=@snapshots  0 0
EOF
    echo -e "${GREEN}Generated /etc/fstab${NC}"
}

# --- Main Execution Flow ---

echo -e "${YELLOW}Gentoo Bootstrap Script${NC}"
echo -e "${YELLOW}Current time: $(date)${NC}"
echo ""

# Init system selection
echo -e "${YELLOW}Select init system:${NC}"
echo "  1) OpenRC (default)"
echo "  2) systemd"
read -r -p "Choice [1]: " init_choice
case "$init_choice" in
    2) INIT_SYSTEM="systemd" ;;
    *) INIT_SYSTEM="openrc" ;;
esac
echo -e "${GREEN}Using init system: ${INIT_SYSTEM}${NC}"

# Step 1: Disk Selection
read -r -p "Ready to select the target disk? (y/n): " confirm_step
if [[ "$confirm_step" != "y" ]]; then
    echo -e "${RED}Aborting.${NC}"
    exit 1
fi
if ! select_target_disk; then
    exit 1
fi

# Step 2: Network Setup
echo -e "${YELLOW}--- MANUAL STEP: Network Setup ---${NC}"
echo "Ensure network connectivity before proceeding."
echo "Use 'net-setup' or manual commands (wpa_supplicant, dhcpcd, etc.)."
echo "Test with: ping -c 3 gentoo.org"
read -r -p "Press Enter once network is configured..."

echo -e "${YELLOW}Testing network...${NC}"
if ping -c 2 gentoo.org >/dev/null 2>&1; then
    echo -e "${GREEN}Network connectivity confirmed.${NC}"
else
    echo -e "${RED}Warning: Cannot reach gentoo.org.${NC}"
    read -r -p "Continue anyway? (y/n): " net_confirm
    if [[ "$net_confirm" != "y" ]]; then
        exit 1
    fi
fi

# Step 3: Partitioning
read -r -p "Ready to partition ${TARGET_DISK_GLOBAL}? DESTRUCTIVE! (y/n): " confirm_step
if [[ "$confirm_step" != "y" ]]; then
    exit 1
fi
if ! partition_and_format_disk; then
    exit 1
fi

# Step 4: Mount Partitions
read -r -p "Ready to mount partitions? (y/n): " confirm_step
if [[ "$confirm_step" != "y" ]]; then
    exit 1
fi
if ! mount_partitions; then
    exit 1
fi

# Step 5: Download & Extract Stage 3
read -r -p "Ready to download and extract Stage 3? (y/n): " confirm_step
if [[ "$confirm_step" != "y" ]]; then
    exit 1
fi
if ! download_and_extract_stage3; then
    exit 1
fi

# Step 6: Prepare Chroot
read -r -p "Ready to prepare chroot environment? (y/n): " confirm_step
if [[ "$confirm_step" != "y" ]]; then
    exit 1
fi
if ! prepare_chroot; then
    exit 1
fi

# Step 7: Copy Portage Configuration
echo -e "${YELLOW}Copying Portage configuration files...${NC}"
if ! copy_portage_config; then
    echo -e "${RED}Warning: Some configuration files may not have been copied.${NC}"
fi

# Step 8: Generate crypttab and fstab
echo -e "${YELLOW}Generating filesystem configuration...${NC}"
generate_crypttab
generate_fstab

echo -e "${GREEN}#####################################################################${NC}"
echo -e "${GREEN}### Initial setup complete!                                       ###${NC}"
echo -e "${GREEN}#####################################################################${NC}"
echo ""
echo -e "${YELLOW}Next steps:${NC}"
echo -e "  1. Run change-root.sh (or enter chroot manually):"
echo -e "     ${GREEN}./change-root.sh${NC}"
echo -e "     Or manually:"
echo -e "     ${GREEN}chroot ${GENTOO_ROOT_MOUNT_GLOBAL} /bin/bash${NC}"
echo -e "     ${GREEN}source /etc/profile${NC}"
echo ""
echo -e "  2. Sync Portage and select profile:"
echo -e "     ${GREEN}emerge --sync${NC}"
echo -e "     ${GREEN}eselect profile list${NC}"
echo -e "     ${GREEN}eselect profile set <number>${NC}"
echo ""
echo -e "${YELLOW}Configuration files copied:${NC}"
echo -e "  - make.conf"
echo -e "  - ccache.conf"
echo -e "  - package.accept_keywords/"
echo -e "  - package.use/"
echo -e "  - sets/"
echo -e "  - dracut.conf.d/ (LUKS + NVIDIA)"
echo ""
echo -e "${YELLOW}Generated:${NC}"
echo -e "  - /etc/crypttab (LUKS root)"
echo -e "  - /etc/fstab (EFI, encrypted swap, Btrfs on LUKS)"
echo ""
echo -e "${YELLOW}LUKS Encryption:${NC}"
echo -e "  Partition: ${PART3_GLOBAL}"
echo -e "  Mapper: ${CRYPT_ROOT_GLOBAL}"
echo -e "  Swap: Encrypted with random key at each boot"
echo ""
echo -e "${YELLOW}Btrfs subvolume layout (on LUKS):${NC}"
echo -e "  @ -> /              (root)"
echo -e "  @home -> /home      (user data)"
echo -e "  @var -> /var        (variable data)"
echo -e "  @log -> /var/log    (logs)"
echo -e "  @snapshots -> /.snapshots (for snapper/timeshift)"
echo ""
echo -e "${YELLOW}Current time: $(date)${NC}"

exit 0