#!/usr/bin/env bash
# btrfs.sh - Btrfs-specific functions for archangel installer
# Source this file after common.sh, config.sh, disk.sh

#############################
# Btrfs Constants
#############################

# Mount options for btrfs subvolumes
BTRFS_OPTS="noatime,compress=zstd,space_cache=v2,discard=async"

# Subvolume layout (matches ZFS dataset structure)
# Format: "name:mountpoint:extra_opts"
BTRFS_SUBVOLS=(
    "@:/::"
    "@home:/home::"
    "@snapshots:/.snapshots::"
    "@var_log:/var/log::"
    "@var_cache:/var/cache::"
    "@tmp:/tmp::nosuid,nodev"
    "@var_tmp:/var/tmp::nosuid,nodev"
    "@media:/media::compress=no"
    "@vms:/vms::nodatacow,compress=no"
    "@var_lib_docker:/var/lib/docker::"
)

#############################
# Btrfs Pre-flight
#############################

btrfs_preflight() {
    step "Checking Btrfs Requirements"

    # Check for btrfs-progs
    if ! command_exists mkfs.btrfs; then
        error "btrfs-progs not installed. Cannot create btrfs filesystem."
    fi
    info "btrfs-progs available."

    # Check for required tools
    require_command btrfs
    require_command grub-install

    info "Btrfs preflight checks passed."
}

#############################
# Btrfs Volume Creation
#############################

create_btrfs_volume() {
    local partition="$1"

    step "Creating Btrfs Filesystem"

    info "Formatting $partition as btrfs..."
    mkfs.btrfs -f -L "archroot" "$partition" || error "Failed to create btrfs filesystem"

    info "Btrfs filesystem created on $partition"
}

#############################
# Subvolume Creation
#############################

create_btrfs_subvolumes() {
    local partition="$1"

    step "Creating Btrfs Subvolumes"

    # Mount the raw btrfs volume temporarily
    mount "$partition" /mnt || error "Failed to mount btrfs volume"

    # Create each subvolume
    for subvol_spec in "${BTRFS_SUBVOLS[@]}"; do
        IFS=':' read -r name mountpoint extra <<< "$subvol_spec"
        info "Creating subvolume: $name -> $mountpoint"
        btrfs subvolume create "/mnt/$name" || error "Failed to create subvolume $name"
    done

    # Unmount raw volume
    umount /mnt

    info "Created ${#BTRFS_SUBVOLS[@]} subvolumes."
}

#############################
# Btrfs Mount Functions
#############################

mount_btrfs_subvolumes() {
    local partition="$1"

    step "Mounting Btrfs Subvolumes"

    # Mount root subvolume first
    info "Mounting @ -> /mnt"
    mount -o "subvol=@,$BTRFS_OPTS" "$partition" /mnt || error "Failed to mount root subvolume"

    # Create mount points and mount remaining subvolumes
    for subvol_spec in "${BTRFS_SUBVOLS[@]}"; do
        IFS=':' read -r name mountpoint extra <<< "$subvol_spec"

        # Skip root, already mounted
        [[ "$name" == "@" ]] && continue

        # Build mount options
        local opts="subvol=$name,$BTRFS_OPTS"

        # Apply extra options (override defaults where specified)
        if [[ -n "$extra" ]]; then
            # Handle compress=no by removing compress from opts and not adding it
            if [[ "$extra" == *"compress=no"* ]]; then
                opts=$(echo "$opts" | sed 's/,compress=zstd//')
            fi
            # Handle nodatacow
            if [[ "$extra" == *"nodatacow"* ]]; then
                opts="$opts,nodatacow"
                opts=$(echo "$opts" | sed 's/,compress=zstd//')
            fi
            # Handle nosuid,nodev for tmp
            if [[ "$extra" == *"nosuid"* ]]; then
                opts="$opts,nosuid,nodev"
            fi
        fi

        info "Mounting $name -> /mnt$mountpoint"
        mkdir -p "/mnt$mountpoint"
        mount -o "$opts" "$partition" "/mnt$mountpoint" || error "Failed to mount $name"
    done

    # Set permissions on tmp directories
    chmod 1777 /mnt/tmp /mnt/var/tmp

    info "All subvolumes mounted."
}

#############################
# Fstab Generation
#############################

generate_btrfs_fstab() {
    local partition="$1"
    local efi_partition="$2"

    step "Generating fstab"

    local uuid
    uuid=$(blkid -s UUID -o value "$partition")

    # Start with header
    cat > /mnt/etc/fstab << EOF
# /etc/fstab - Btrfs subvolume mounts
# IMPORTANT: Using subvol= NOT subvolid= for snapshot compatibility
# Generated by archangel installer

EOF

    # Add each subvolume
    for subvol_spec in "${BTRFS_SUBVOLS[@]}"; do
        IFS=':' read -r name mountpoint extra <<< "$subvol_spec"

        # Build mount options
        local opts="subvol=$name,$BTRFS_OPTS"

        # Apply extra options
        if [[ -n "$extra" ]]; then
            if [[ "$extra" == *"compress=no"* ]]; then
                opts=$(echo "$opts" | sed 's/,compress=zstd//')
            fi
            if [[ "$extra" == *"nodatacow"* ]]; then
                opts="$opts,nodatacow"
                opts=$(echo "$opts" | sed 's/,compress=zstd//')
            fi
            if [[ "$extra" == *"nosuid"* ]]; then
                opts="$opts,nosuid,nodev"
            fi
        fi

        echo "UUID=$uuid  $mountpoint  btrfs  $opts  0 0" >> /mnt/etc/fstab
    done

    # Add EFI partition
    local efi_uuid
    efi_uuid=$(blkid -s UUID -o value "$efi_partition")
    echo "" >> /mnt/etc/fstab
    echo "# EFI System Partition" >> /mnt/etc/fstab
    echo "UUID=$efi_uuid  /efi  vfat  defaults,noatime  0 2" >> /mnt/etc/fstab

    info "fstab generated with ${#BTRFS_SUBVOLS[@]} btrfs mounts + EFI"
}

#############################
# Snapper Configuration
#############################

configure_snapper() {
    step "Configuring Snapper"

    # Create snapper config directory
    mkdir -p /mnt/etc/snapper/configs

    # Snapper config for root - create manually (avoids D-Bus requirement in chroot)
    # Note: We already have @snapshots subvolume mounted at /.snapshots
    info "Creating snapper config for root..."

    # Set snapper timeline settings
    # Keep: 6 hourly, 7 daily, 2 weekly, 1 monthly
    info "Configuring snapshot retention policy..."
    cat > /mnt/etc/snapper/configs/root << 'EOF'
# Snapper config for root filesystem
SUBVOLUME="/"
FSTYPE="btrfs"
QGROUP=""

# Automatic timeline snapshots
TIMELINE_CREATE="yes"
TIMELINE_CLEANUP="yes"

# Retention policy
TIMELINE_MIN_AGE="1800"
TIMELINE_LIMIT_HOURLY="6"
TIMELINE_LIMIT_DAILY="7"
TIMELINE_LIMIT_WEEKLY="2"
TIMELINE_LIMIT_MONTHLY="1"
TIMELINE_LIMIT_YEARLY="0"

# Cleanup settings
NUMBER_CLEANUP="yes"
NUMBER_MIN_AGE="1800"
NUMBER_LIMIT="50"
NUMBER_LIMIT_IMPORTANT="10"

# User access
ALLOW_USERS=""
ALLOW_GROUPS=""

# Sync ACL
SYNC_ACL="no"

# Background comparison
BACKGROUND_COMPARISON="yes"

# Empty pre-post cleanup
EMPTY_PRE_POST_CLEANUP="yes"
EMPTY_PRE_POST_MIN_AGE="1800"
EOF

    # Register the config with snapper
    mkdir -p /mnt/etc/sysconfig
    echo 'SNAPPER_CONFIGS="root"' > /mnt/etc/sysconfig/snapper

    # Enable snapper timers
    arch-chroot /mnt systemctl enable snapper-timeline.timer
    arch-chroot /mnt systemctl enable snapper-cleanup.timer

    info "Snapper configured with timeline snapshots enabled."
}

#############################
# GRUB Configuration
#############################

configure_grub() {
    local efi_partition="$1"

    step "Configuring GRUB Bootloader"

    # Mount EFI partition
    mkdir -p /mnt/efi
    mount "$efi_partition" /mnt/efi

    # Configure GRUB defaults
    info "Setting GRUB configuration..."
    cat > /mnt/etc/default/grub << 'EOF'
# GRUB configuration for btrfs root with snapshots
GRUB_DEFAULT=0
GRUB_TIMEOUT=5
GRUB_DISTRIBUTOR="Arch"
GRUB_CMDLINE_LINUX_DEFAULT="loglevel=3 quiet"
GRUB_CMDLINE_LINUX=""

# Disable os-prober (single-boot system)
GRUB_DISABLE_OS_PROBER=true
EOF

    # Create /boot/grub directory (grub-install expects this)
    mkdir -p /mnt/boot/grub

    # Install GRUB to EFI
    info "Installing GRUB to EFI partition..."
    arch-chroot /mnt grub-install --target=x86_64-efi --efi-directory=/efi --bootloader-id=GRUB \
        || error "GRUB installation failed"

    # Generate GRUB config (without grub-btrfs first time)
    info "Generating GRUB configuration..."
    arch-chroot /mnt grub-mkconfig -o /boot/grub/grub.cfg \
        || error "Failed to generate GRUB config"

    # Enable grub-btrfsd for automatic snapshot menu updates
    info "Enabling grub-btrfs daemon..."
    arch-chroot /mnt systemctl enable grub-btrfsd

    info "GRUB configured with btrfs snapshot support."
}

#############################
# Pacman Snapshot Hook
#############################

configure_btrfs_pacman_hook() {
    step "Configuring Pacman Snapshot Hook"

    # snap-pac handles this automatically when installed
    # Just verify it's set up
    info "snap-pac will create pre/post snapshots for pacman transactions."
    info "Snapshots visible in GRUB menu via grub-btrfs."
}

#############################
# Genesis Snapshot
#############################

create_btrfs_genesis_snapshot() {
    step "Creating Genesis Snapshot"

    # Create snapshot manually (snapper requires D-Bus which isn't available in chroot)
    # Snapper stores snapshots in /.snapshots/<num>/snapshot as read-only subvolumes
    local snap_dir="/mnt/.snapshots/1"
    mkdir -p "$snap_dir"

    # Create the snapshot of root
    btrfs subvolume snapshot -r /mnt "$snap_dir/snapshot" \
        || error "Failed to create genesis snapshot"

    # Create snapper info.xml (required for snapper to recognize the snapshot)
    local timestamp
    timestamp=$(date -u +"%Y-%m-%d %H:%M:%S")
    cat > "$snap_dir/info.xml" << EOF
<?xml version="1.0"?>
<snapshot>
  <type>single</type>
  <num>1</num>
  <date>$timestamp</date>
  <description>genesis</description>
  <cleanup>number</cleanup>
</snapshot>
EOF

    info "Genesis snapshot created at /.snapshots/1/snapshot"
    info "Restore with: snapper -c root rollback 1"
}

#############################
# Btrfs Services
#############################

configure_btrfs_services() {
    step "Configuring System Services"

    # Enable standard services
    arch-chroot /mnt systemctl enable NetworkManager
    arch-chroot /mnt systemctl enable avahi-daemon

    # Snapper timers (already enabled in configure_snapper)

    # grub-btrfsd (already enabled in configure_grub)

    info "System services configured."
}

#############################
# Btrfs Initramfs
#############################

configure_btrfs_initramfs() {
    step "Configuring Initramfs for Btrfs"

    # Backup original
    cp /mnt/etc/mkinitcpio.conf /mnt/etc/mkinitcpio.conf.bak

    # Remove archiso drop-in if present
    if [[ -f /mnt/etc/mkinitcpio.conf.d/archiso.conf ]]; then
        info "Removing archiso drop-in config..."
        rm -f /mnt/etc/mkinitcpio.conf.d/archiso.conf
    fi

    # Create proper linux-lts preset
    info "Creating linux-lts preset..."
    cat > /mnt/etc/mkinitcpio.d/linux-lts.preset << 'EOF'
# mkinitcpio preset file for linux-lts

PRESETS=(default fallback)

ALL_kver="/boot/vmlinuz-linux-lts"

default_image="/boot/initramfs-linux-lts.img"

fallback_image="/boot/initramfs-linux-lts-fallback.img"
fallback_options="-S autodetect"
EOF

    # Configure hooks for btrfs
    # btrfs module is built into kernel, but we need the btrfs hook for multi-device
    sed -i 's/^HOOKS=.*/HOOKS=(base udev microcode modconf kms keyboard keymap consolefont block filesystems fsck)/' \
        /mnt/etc/mkinitcpio.conf

    # Regenerate initramfs
    info "Regenerating initramfs..."
    arch-chroot /mnt mkinitcpio -P

    info "Initramfs configured for btrfs."
}

#############################
# Btrfs Cleanup
#############################

btrfs_cleanup() {
    step "Cleaning Up Btrfs"

    # Unmount in reverse order
    info "Unmounting subvolumes..."

    # Unmount EFI first
    umount /mnt/efi 2>/dev/null || true

    # Unmount all btrfs subvolumes (reverse order)
    for ((i=${#BTRFS_SUBVOLS[@]}-1; i>=0; i--)); do
        IFS=':' read -r name mountpoint extra <<< "${BTRFS_SUBVOLS[$i]}"
        [[ "$name" == "@" ]] && continue
        umount "/mnt$mountpoint" 2>/dev/null || true
    done

    # Unmount root last
    umount /mnt 2>/dev/null || true

    info "Btrfs cleanup complete."
}