#!/bin/sh
# zfs-pool-migrate.sh — Pool migration wizard
#
# Handles three migration methods:
#   1. Disk-by-disk replacement (4 bays needed)
#   2. New pool + transfer (6 bays needed, or USB enclosures)
#   3. Backup + restore (external storage needed)

set -e

POOL_NAME="clawdie"
LOG="/var/log/clawdie-migrate.log"

. "/usr/local/share/clawdie-iso/build.cfg"

dialog() { bsddialog --backtitle "Clawdie Pool Migration" "$@" ; }
die() { echo "ERROR: $1" >&2; exit 1; }
log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" >> "$LOG"; }

get_disk_info() {
    local disk=$1
    local info=""
    
    info+="DEVICE:  $disk\n"
    
    local model=$(camcontrol inquiry $disk 2>/dev/null | head -1)
    [ -n "$model" ] && info+="MODEL:   $model\n"
    
    local serial=$(camcontrol identify $disk 2>/dev/null | grep -i 'serial number' | cut -d: -f2 | tr -d ' ')
    [ -n "$serial" ] && info+="SERIAL: $serial\n"
    
    local size_bytes=$(diskinfo $disk 2>/dev/null | awk '{print $3}')
    if [ -n "$size_bytes" ]; then
        local size_gb=$((size_bytes / 1073741824))
        info+="SIZE:   ${size_gb} GB\n"
    fi
    
    local temp=$(smartctl -A $disk 2>/dev/null | grep -i temperature | awk '{print $10}')
    [ -n "$temp" ] && info+="TEMP:   ${temp}°C\n"
    
    local health=$(smartctl -H $disk 2>/dev/null | grep -E '^SMART overall-health' | awk '{print $NF}')
    [ -n "$health" ] && info+="HEALTH: $health\n"
    
    local slot=$(sesutil map 2>/dev/null | grep "^$disk" | awk '{print $3}')
    [ -n "$slot" ] && info+="SLOT:   $slot\n"
    
    echo "$info"
}

get_pool_status() {
    local status=$(zpool status $POOL_NAME 2>/dev/null)
    echo "$status"
}

get_pool_usage() {
    local used=$(zfs list -Hp -o used $POOL_NAME 2>/dev/null)
    local avail=$(zfs list -Hp -o avail $POOL_NAME 2>/dev/null)
    local total=$((used + avail))
    local used_gb=$((used / 1073741824))
    local total_gb=$((total / 1073741824))
    local percent=$((used * 100 / total))
    
    echo "USED:   ${used_gb} GB / ${total_gb} GB (${percent}%)"
}

count_disks_in_pool() {
    zpool status $POOL_NAME 2>/dev/null | grep -E '^\s+(ada|da|nvd|nda)' | wc -l
}

get_pool_disks() {
    zpool status $POOL_NAME 2>/dev/null | grep -E '^\s+(ada|da|nvd|nda)' | awk '{print $1}'
}

get_available_disks() {
    local pool_disks=$(get_pool_disks | sort -u)
    local all_disks=$(camcontrol devlist 2>/dev/null | grep -oE '(ada|da|nvd|nda)' | awk '{print $NF}' | sort -u)
    
    for d in $all_disks; do
        if ! echo "$pool_disks" | grep -qw "$d"; then
            echo "$d"
        fi
    done
}

format_size() {
    local bytes=$1
    if [ $bytes -ge 1099511627776 ]; then
        echo "$((bytes / 1099511627776)) TB"
    elif [ $bytes -ge 1073741824 ]; then
        echo "$((bytes / 1073741824)) GB"
    elif [ $bytes -ge 1048576 ]; then
        echo "$((bytes / 1048576)) MB"
    else
        echo "$bytes bytes"
    fi
}

method_disk_by_disk() {
    log "Starting disk-by-disk migration"
    
    local pool_disks=$(get_pool_disks)
    local disk_count=$(echo "$pool_disks" | wc -l)
    local available=$(get_available_disks)
    
    if [ -z "$available" ]; then
        dialog --msgbox \
            "\nNo replacement disks available.\n\nConnect replacement disks and restart." \
            10 50
        return 1
    fi
    
    local step=1
    local total_steps=$disk_count
    
    while [ $step -le $total_steps ]; do
        local current_disk=$(echo "$pool_disks" | sed -n "${step}p")
        local disk_info=$(get_disk_info $current_disk)
        
        local replacement=$(dialog --menu \
            "Step ${step} of ${total_steps}: Replace ${current_disk}\n\n${disk_info}\n\nSelect replacement disk:" \
            20 70 10 \
            $(echo "$available" | while read -r disk; do echo "$disk \"$(get_disk_info $disk | head -1)\""; done) \
            3>&1 1>&2 2>&3)
        
        [ $? -ne 0 ] && return 1
        
        local replacement_size=$(diskinfo $replacement 2>/dev/null | awk '{print $3}')
        local current_size=$(diskinfo $current_disk 2>/dev/null | awk '{print $3}')
        
        if [ $replacement_size -lt $current_size ]; then
            dialog --msgbox \
                "\nReplacement disk is smaller than current disk.\n\nReplacement: $(format_size $replacement_size)\nCurrent: $(format_size $current_size)\n\nUse a disk of equal or larger size." \
                12 60
            continue
        fi
        
        local confirm=$(dialog --yesno \
            "\nConfirm replacement:\n\n  ${current_disk} → ${replacement}\n\nThis will take 4-8 hours.\nPool remains usable during resilver.\n\nContinue?" \
            12 60)
        
        [ "$confirm" != "yes" ] && continue
        
        log "Replacing $current_disk with $replacement"
        
        dialog --infobox "Starting resilver...\nThis screen will update periodically." 6 50
        
        zpool replace $POOL_NAME $current_disk $replacement 2>&1 | tee -a "$LOG"
        
        local progress=0
        while true; do
            local status=$(zpool status $POOL_NAME 2>/dev/null)
            if echo "$status" | grep -q "resilver completed"; then
                break
            fi
            
            local scan_line=$(echo "$status" | grep -A1 "scan:")
            if echo "$scan_line" | grep -qE '[0-9]+%'; then
                progress=$(echo "$scan_line" | grep -oE '[0-9]+' | head -1 | tr -d '%')
                dialog --infobox "Resilver in progress: ${progress}%\n\nReplacing: ${current_disk} → ${replacement}\n\nDo not power off." 6 50
            fi
            
            sleep 30
        done
        
        log "Resilver completed for $current_disk → $replacement"
        
        pool_disks=$(echo "$pool_disks" | sed "s/${current_disk}/${replacement}/")
        
        step=$((step + 1))
    done
    
    dialog --msgbox \
        "\nDisk-by-disk migration complete!\n\nAll disks have been replaced.\nPool remains redundant throughout the process.\n\nConsider running:\n  zpool scrub $POOL_NAME\n\nto verify data integrity." \
        15 60
    
    log "Disk-by-disk migration completed successfully"
    return 0
}

method_new_pool_transfer() {
    log "Starting new pool + transfer migration"
    
    local available=$(get_available_disks | sort -u)
    local available_count=$(echo "$available" | wc -l)
    
    if [ $available_count -lt 3 ]; then
        local use_usb=$(dialog --yesno \
            "\nNot enough internal disks for new pool.\n\nFound: ${available_count} available\nNeeded: 3 disks\n\nUse USB enclosures for additional disks?" \
            12 60)
        
        if [ "$use_usb" != "yes" ]; then
            dialog --msgbox \
                "\nCannot proceed.\n\nConnect additional disks (USB enclosures work) and restart.\n\nAlternatively, use disk-by-disk method (4 bays) or backup/restore (external storage)." \
                12 60
            return 1
        fi
    fi
    
    local selected=$(dialog --checklist \
        "Select 3 disks for new migration pool:\n\nThese disks will be ERASED." \
        20 70 10 \
        $(echo "$available" | while read -r disk; do echo "$disk \"$(get_disk_info $disk | grep MODEL | cut -d: -f2-)\" off"; done) \
        3>&1 1>&2 2>&3)
    
    [ $? -ne 0 ] && return 1
    
    local selected_count=$(echo "$selected" | wc -w)
    if [ $selected_count -ne 3 ]; then
        dialog --msgbox "\nPlease select exactly 3 disks.\n\nSelected: ${selected_count}" 10 50
        return 1
    fi
    
    local raid_type=$(dialog --radiolist \
        "RAID level for new pool:" 12 60 3 \
        "raidz1" "3 disks, 1 fault tolerance (recommended)" on \
        "raidz2" "4+ disks, 2 fault tolerance" off \
        "mirror" "2 disks, 1 fault tolerance, faster reads" off \
        3>&1 1>&2 2>&3)
    
    [ $? -ne 0 ] && return 1
    
    local confirm=$(dialog --yesno \
        "\nCreate new pool and transfer data?\n\nNew pool: clawdie-new\nRAID: ${raid_type}\nDisks: ${selected}\n\n⚠ All data on selected disks will be DESTROYED.\n\nTransfer time: ~4-8 hours" \
        15 70)
    
    [ "$confirm" != "yes" ] && return 1
    
    log "Creating new pool clawdie-new with disks: $selected"
    
    local disk_args=""
    for disk in $selected; do
        disk_args="$disk_args /dev/$disk"
    done
    
    zpool create clawdie-new $raid_type $disk_args 2>&1 | tee -a "$LOG"
    
    [ $? -ne 0 ] && die "Failed to create new pool"
    
    dialog --infobox "Transferring data...\nThis will take 4-8 hours.\nProgress shown periodically." 6 60
    
    log "Starting zfs send/recv from $POOL_NAME to clawdie-new"
    
    local progress_file="/tmp/zfs-transfer-progress"
    
    zfs snapshot -r $POOL_NAME@migration-transfer 2>&1 | tee -a "$LOG"
    
    ( zfs send -Rv $POOL_NAME@migration-transfer 2>&1 | tee "$progress_file" | \
      pv -l -F -p -b -t -r -i > "$progress_file.stats" 2>/dev/null | \
      zfs recv -v clawdie-new 2>&1 | tee -a "$LOG" ) || \
    zfs send -R $POOL_NAME@migration-transfer 2>&1 | zfs recv -F clawdie-new 2>&1 | tee -a "$LOG"
    
    [ $? -ne 0 ] && die "Transfer failed"
    
    log "Transfer completed"
    
    local verify=$(dialog --yesno \
        "\nData transfer complete!\n\nVerify data integrity? (recommended)\nThis compares source and destination." \
        10 60)
    
    if [ "$verify" = "yes" ]; then
        dialog --infobox "Verifying data integrity..." 5 50
        zfs diff $POOL_NAME@migration-transfer clawdie-new 2>&1 | tee -a "$LOG"
    fi
    
    local destroy_old=$(dialog --yesno \
        "\nMigration complete!\n\nOld pool '$POOL_NAME' is still intact.\n\nDestroy old pool now?\n(Recommended: keep old disks as backup for 1-2 weeks)" \
        12 70)
    
    if [ "$destroy_old" = "yes" ]; then
        log "Destroying old pool $POOL_NAME"
        zpool export $POOL_NAME 2>&1 | tee -a "$LOG"
        dialog --msgbox \
            "\nOld pool exported.\n\nOld disks can now be removed.\n\nNew pool 'clawdie-new' needs to be renamed:\n\n  zpool export clawdie-new\n  zpool import clawdie-new clawdie" \
            12 60
    else
        dialog --msgbox \
            "\nBoth pools exist:\n  - $POOL_NAME (old)\n  - clawdie-new (new)\n\nTo complete migration:\n  1. Boot from USB\n  2. Export old pool\n  3. Rename new pool:\n     zpool export clawdie-new\n     zpool import clawdie-new clawdie\n  4. Reboot" \
            15 70
    fi
    
    log "New pool + transfer migration completed"
    return 0
}

method_backup_restore() {
    log "Starting backup + restore migration"
    
    local backup_targets=""
    
    backup_targets+="local External USB drive\n"
    backup_targets+="network Network (SSH/SCP)\n"
    backup_targets+="zfs Another ZFS pool\n"
    
    local dest_type=$(dialog --menu \
        "Select backup destination:" 12 60 3 \
        $backup_targets \
        3>&1 1>&2 2>&3)
    
    [ $? -ne 0 ] && return 1
    
    local dest_path=""
    local dest_size=0
    
    case "$dest_type" in
        "local")
            local usb_disks=$(camcontrol devlist 2>/dev/null | grep -oE '(da|ada)[0-9]+' | sort -u)
            
            if [ -z "$usb_disks" ]; then
                dialog --msgbox "\nNo external drives detected.\n\nConnect external drive and restart." 10 50
                return 1
            fi
            
            dest_path=$(dialog --menu \
                "Select external drive:" 12 60 5 \
                $(echo "$usb_disks" | while read -r disk; do echo "$disk \"$(get_disk_info $disk | grep MODEL | cut -d: -f2-)\""; done) \
                3>&1 1>&2 2>&3)
            
            [ $? -ne 0 ] && return 1
            
            dest_path="/dev/$dest_path"
            dest_size=$(diskinfo $dest_path 2>/dev/null | awk '{print $3}')
            ;;
            
        "network")
            dest_path=$(dialog --inputbox \
                "Network destination (user@host:/path):" 8 60 \
                "user@backup-server:/backups/clawdie" \
                3>&1 1>&2 2>&3)
            [ $? -ne 0 ] && return 1
            ;;
            
        "zfs")
            local other_pools=$(zpool import 2>/dev/null | grep "pool:" | awk '{print $2}' | grep -v "^$POOL_NAME$")
            
            if [ -z "$other_pools" ]; then
                dialog --msgbox "\nNo other pools available.\n\nCreate a backup pool first." 10 50
                return 1
            fi
            
            dest_path=$(dialog --menu \
                "Select backup pool:" 10 60 3 \
                $(echo "$other_pools" | while read -r pool; do echo "$pool \"ZFS pool\""; done) \
                3>&1 1>&2 2>&3)
            [ $? -ne 0 ] && return 1
            ;;
    esac
    
    local pool_used=$(zfs list -Hp -o used $POOL_NAME 2>/dev/null)
    local pool_used_gb=$((pool_used / 1073741824))
    
    local confirm=$(dialog --yesno \
        "\nBackup configuration:\n\n  Source: $POOL_NAME (${pool_used_gb} GB)\n  Destination: $dest_path\n\nTime: ~4-8 hours\n\nStart backup?" \
        12 60)
    
    [ "$confirm" != "yes" ] && return 1
    
    log "Starting backup to $dest_path"
    
    dialog --infobox "Creating backup...\nThis will take 4-8 hours." 6 50
    
    case "$dest_type" in
        "local")
            local mount_point="/mnt/backup"
            mkdir -p "$mount_point"
            newfs "$dest_path" 2>&1 | tee -a "$LOG" || mount -t ufs "$dest_path" "$mount_point" 2>&1 | tee -a "$LOG"
            zfs send -R $POOL_NAME@backup-$(date +%Y.%m.%d) > "$mount_point/clawdie-backup.zfs" 2>&1 | tee -a "$LOG"
            umount "$mount_point"
            ;;
        "network")
            zfs send -R $POOL_NAME@backup-$(date +%Y.%m.%d) | ssh "$dest_path" "cat > clawdie-backup.zfs" 2>&1 | tee -a "$LOG"
            ;;
        "zfs")
            zfs send -R $POOL_NAME@backup-$(date +%Y.%m.%d) | zfs recv "${dest_path}/clawdie-backup" 2>&1 | tee -a "$LOG"
            ;;
    esac
    
    [ $? -ne 0 ] && die "Backup failed"
    
    log "Backup completed"
    
    dialog --msgbox \
        "\nBackup complete!\n\nNow:\n  1. Power off\n  2. Replace old disks with new disks\n  3. Boot USB again\n  4. Select 'Restore from backup'\n\nOld pool will be destroyed." \
        12 60
    
    log "Backup + restore migration - backup phase completed"
    return 0
}

main() {
    log "Pool migration wizard started"
    
    kldload zfs 2>/dev/null || true
    
    if ! zpool list $POOL_NAME >/dev/null 2>&1; then
        dialog --msgbox \
            "\nPool '$POOL_NAME' not found.\n\nMigration requires an existing pool." \
            10 50
        return 1
    fi
    
    local pool_status=$(get_pool_status)
    local pool_usage=$(get_pool_usage)
    local disk_count=$(count_disks_in_pool)
    
    local method=$(dialog --radiolist \
        "Pool Migration\n\nPool: $POOL_NAME\n${pool_usage}\nDisks: ${disk_count}\n\nSelect migration method:" \
        18 70 3 \
        "disk-by-disk" "Replace one disk at a time (4 bays)" on \
        "new-pool" "Create new pool + transfer (6 bays)" off \
        "backup-restore" "Backup + restore (external storage)" off \
        3>&1 1>&2 2>&3)
    
    [ $? -ne 0 ] && return 1
    
    case "$method" in
        "disk-by-disk")
            method_disk_by_disk
            ;;
        "new-pool")
            method_new_pool_transfer
            ;;
        "backup-restore")
            method_backup_restore
            ;;
        *)
            die "Unknown method: $method"
            ;;
    esac
}

main "$@"