#!/bin/bash
# =============================================================================
# restore_windows.sh
# Windows 11 NVMe Restore Script
# Backup-Verzeichnis = Verzeichnis in dem das Script liegt
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BACKUP_DIR="$SCRIPT_DIR"

# Farben
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

# Quell-NVMe Größe in Sektoren (aus windows11.partitions)
SOURCE_SECTORS=1000215182
SOURCE_SIZE_GB=476

# =============================================================================
# Funktionen
# =============================================================================

print_header() {
    echo ""
    echo -e "${BOLD}${BLUE}=============================================${NC}"
    echo -e "${BOLD}${BLUE}   Windows 11 NVMe Restore Script${NC}"
    echo -e "${BOLD}${BLUE}=============================================${NC}"
    echo ""
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        echo -e "${RED}[FEHLER] Dieses Script muss als root ausgeführt werden!${NC}"
        exit 1
    fi
}

check_backups() {
    echo -e "${BOLD}[1/4] Prüfe Backup-Dateien in: ${CYAN}$BACKUP_DIR${NC}"
    echo ""

    local missing=0
    local files=("efi.img" "msr.img" "recovery.ntfsclone.gz" "windows.ntfsclone" "windows11.partitions")

    for f in "${files[@]}"; do
        if [[ -f "$BACKUP_DIR/$f" ]]; then
            local size
            size=$(du -sh "$BACKUP_DIR/$f" | cut -f1)
            echo -e "  ${GREEN}✓${NC} $f ${CYAN}($size)${NC}"
        else
            echo -e "  ${RED}✗${NC} $f ${RED}(FEHLT!)${NC}"
            missing=$((missing + 1))
        fi
    done

    echo ""
    if [[ $missing -gt 0 ]]; then
        echo -e "${RED}[FEHLER] $missing Backup-Datei(en) fehlen in $BACKUP_DIR${NC}"
        exit 1
    fi
    echo -e "${GREEN}Alle Backup-Dateien vorhanden.${NC}"
    echo ""
}

detect_usb_devices() {
    echo -e "${BOLD}[2/4] Suche USB-SSDs...${NC}"
    echo ""

    # Nur USB-Geräte, keine Loop-Devices, nur Top-Level Devices
    mapfile -t USB_DEVICES < <(lsblk -d -o NAME,SIZE,MODEL,TRAN | awk 'NR>1 && $4=="usb" {print $1}')

    if [[ ${#USB_DEVICES[@]} -eq 0 ]]; then
        echo -e "${RED}[FEHLER] Keine USB-Geräte gefunden!${NC}"
        exit 1
    fi

    echo -e "  Gefundene USB-Geräte:"
    echo ""

    declare -g -A DEVICE_SIZES
    declare -g -a DEVICE_LIST=()

    for dev in "${USB_DEVICES[@]}"; do
        local info
        info=$(lsblk -d -o NAME,SIZE,MODEL "/dev/$dev" | tail -1)
        local size_bytes
        size_bytes=$(lsblk -d -b -o SIZE "/dev/$dev" | tail -1 | tr -d ' ')
        local size_gb=$(( size_bytes / 1024 / 1024 / 1024 ))
        DEVICE_SIZES[$dev]=$size_gb

        echo -e "  ${CYAN}[$((${#DEVICE_LIST[@]} + 1))]${NC} /dev/$dev  →  $info"
        DEVICE_LIST+=("$dev")
    done

    echo ""
}

select_device() {
    echo -e "${BOLD}[3/4] Ziel-Device auswählen${NC}"
    echo ""
    echo -e "  Welches Device soll bespielt werden?"
    echo -e "  ${YELLOW}⚠  ACHTUNG: Alle Daten auf dem Ziel werden gelöscht!${NC}"
    echo ""

    local choice
    while true; do
        read -rp "  Auswahl (1-${#DEVICE_LIST[@]}): " choice
        if [[ "$choice" =~ ^[0-9]+$ ]] && [[ "$choice" -ge 1 ]] && [[ "$choice" -le ${#DEVICE_LIST[@]} ]]; then
            TARGET_DEV="${DEVICE_LIST[$((choice - 1))]}"
            TARGET_SIZE_GB="${DEVICE_SIZES[$TARGET_DEV]}"
            echo ""
            echo -e "  Gewählt: ${CYAN}/dev/$TARGET_DEV${NC} (${TARGET_SIZE_GB} GB)"
            break
        else
            echo -e "  ${RED}Ungültige Eingabe.${NC}"
        fi
    done
    echo ""
}

determine_mode() {
    echo -e "${BOLD}[4/4] Größenvergleich${NC}"
    echo ""
    echo -e "  Quelle:  ${SOURCE_SIZE_GB} GB"
    echo -e "  Ziel:    ${TARGET_SIZE_GB} GB"
    echo ""

    if [[ $TARGET_SIZE_GB -lt 100 ]]; then
        echo -e "${RED}[FEHLER] Ziel-Device ist zu klein (< 100 GB). Mindestens 110 GB benötigt.${NC}"
        exit 1
    fi

    if [[ $TARGET_SIZE_GB -ge $((SOURCE_SIZE_GB - 10)) ]]; then
        MODE="direct"
        echo -e "  ${GREEN}✓ Ziel ist groß genug → Direktes Restore möglich${NC}"
    else
        MODE="shrink"
        echo -e "  ${YELLOW}⚠ Ziel ist kleiner → Windows-Partition wird zuerst verkleinert${NC}"
        # Ziel-Größe für ntfsresize: 6 GB Puffer zur Partition
        SHRINK_GB=$(( TARGET_SIZE_GB - 8 ))
        echo -e "  Verkleinere Windows auf: ${SHRINK_GB} GB"
    fi
    echo ""
}

confirm() {
    echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${BOLD}  Zusammenfassung:${NC}"
    echo -e "  Backup-Quelle:  $BACKUP_DIR"
    echo -e "  Ziel-Device:    /dev/$TARGET_DEV (${TARGET_SIZE_GB} GB)"
    echo -e "  Modus:          $MODE"
    echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo ""
    echo -e "${RED}${BOLD}  ALLE DATEN AUF /dev/$TARGET_DEV WERDEN GELÖSCHT!${NC}"
    echo ""
    read -rp "  Fortfahren? (ja/nein): " confirm
    if [[ "$confirm" != "ja" ]]; then
        echo -e "${YELLOW}Abgebrochen.${NC}"
        exit 0
    fi
    echo ""
}

write_partition_table() {
    echo -e "${BOLD}▶ Schreibe Partitionstabelle...${NC}"

    # Sektorzahl des Ziel-Devices ermitteln
    local target_sectors
    target_sectors=$(lsblk -d -b -o SIZE "/dev/$TARGET_DEV" | tail -1 | tr -d ' ')
    target_sectors=$(( target_sectors / 512 ))

    # GPT reserviert die letzten 33 Sektoren für die sekundäre Partitionstabelle
    local last_lba=$(( target_sectors - 34 ))

    # Recovery-Partition ans Ende setzen (2 GB = 4096000 Sektoren)
    local recovery_size=4096000
    local recovery_start=$(( last_lba - recovery_size + 1 ))
    local windows_size=$(( recovery_start - 567296 ))

    umount "/dev/${TARGET_DEV}"* 2>/dev/null || true

    sfdisk "/dev/$TARGET_DEV" << EOF
label: gpt
label-id: 97DAE1BC-08B3-42DA-8F44-B1F2511FD833
unit: sectors
first-lba: 34
sector-size: 512
/dev/${TARGET_DEV}1 : start=2048,	size=532480,        type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B, uuid=C50622FB-E1FE-4D65-9F3B-4DC0B617F>
/dev/${TARGET_DEV}2 : start=534528,     size=32768,         type=E3C9E316-0B5C-4DB8-817D-F92DF00215AE, uuid=404C486F-D517-4FD2-9F5C-CDF63678F>
/dev/${TARGET_DEV}3 : start=567296,     size=${windows_size}, type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7, uuid=5CB307A8-5FBE-4F89-B7A9-43602C8>
/dev/${TARGET_DEV}4 : start=${recovery_start}, size=${recovery_size}, type=DE94BBA4-06D1-4D40-A16A-BFD50179D6AC, uuid=CB3833DE-DA71-44BE-8553>
EOF

    echo -e "${GREEN}✓ Partitionstabelle geschrieben${NC}"
    echo ""
}

write_efi_msr() {
    echo -e "${BOLD}▶ Schreibe EFI + MSR...${NC}"
    dd if="$BACKUP_DIR/efi.img" of="/dev/${TARGET_DEV}1" bs=4M status=progress
    dd if="$BACKUP_DIR/msr.img" of="/dev/${TARGET_DEV}2" bs=4M status=progress
    echo -e "${GREEN}✓ EFI + MSR geschrieben${NC}"
    echo ""
}

write_windows_direct() {
    echo -e "${BOLD}▶ Schreibe Windows-Partition (direkt)...${NC}"
    ntfsclone --restore-image -O "/dev/${TARGET_DEV}3" "$BACKUP_DIR/windows.ntfsclone"
    echo -e "${GREEN}✓ Windows-Partition geschrieben${NC}"
    echo ""
}

write_windows_shrink() {
    echo -e "${BOLD}▶ Erstelle temporäres raw Image...${NC}"
    local tmp_img="$BACKUP_DIR/windows_tmp_$$.img"

    # Aufräumen bei Abbruch
    trap "losetup -d /dev/loop4 2>/dev/null; rm -f '$tmp_img'; echo 'Aufgeräumt.'" EXIT

    ntfsclone --restore-image -O "$tmp_img" "$BACKUP_DIR/windows.ntfsclone"

    echo -e "${BOLD}▶ Verkleinere NTFS auf ${SHRINK_GB} GB...${NC}"
    losetup /dev/loop4 "$tmp_img"
    echo "y" | ntfsresize -s "${SHRINK_GB}G" /dev/loop4

    echo -e "${BOLD}▶ Schreibe verkleinerte Windows-Partition...${NC}"
    ntfsclone --force -O "/dev/${TARGET_DEV}3" /dev/loop4

    losetup -d /dev/loop4
    rm -f "$tmp_img"
    trap - EXIT

    echo -e "${GREEN}✓ Windows-Partition geschrieben (verkleinert)${NC}"
    echo ""
}

write_recovery() {
    echo -e "${BOLD}▶ Schreibe Recovery-Partition...${NC}"
    gunzip -c "$BACKUP_DIR/recovery.ntfsclone.gz" | ntfsclone --restore-image -O "/dev/${TARGET_DEV}4" -
    echo -e "${GREEN}✓ Recovery-Partition geschrieben${NC}"
    echo ""
}

resize_windows() {
    echo -e "${BOLD}▶ Erweitere Windows-Partition auf volle Größe...${NC}"
    echo "y" | ntfsresize --force "/dev/${TARGET_DEV}3" || ntfsresize --force "/dev/${TARGET_DEV}3"
    echo -e "${GREEN}✓ NTFS erweitert${NC}"
    echo ""
}

# =============================================================================
# Hauptprogramm
# =============================================================================

print_header
check_root
check_backups
detect_usb_devices
select_device
determine_mode
confirm

write_partition_table
write_efi_msr

if [[ "$MODE" == "direct" ]]; then
    write_windows_direct
else
    write_windows_shrink
fi

write_recovery
resize_windows

echo ""
echo -e "${GREEN}${BOLD}=============================================${NC}"
echo -e "${GREEN}${BOLD}   Restore erfolgreich abgeschlossen! ✓${NC}"
echo -e "${GREEN}${BOLD}=============================================${NC}"
echo ""
echo -e "  Die NVMe kann jetzt in den Laptop eingebaut werden."
echo -e "  Windows führt beim ersten Boot automatisch chkdsk aus."
echo ""