#!/bin/bash
# =============================================================================
# restore_windows.sh
# Windows 11 NVMe Restore Script - Universell für alle Laptop-Modelle
# Backup-Verzeichnis = Verzeichnis in dem das Script liegt
# Alle Werte werden automatisch aus windows11.partitions gelesen
# =============================================================================

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'

# =============================================================================
# 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 ""
}

read_partition_info() {
    # Alle Werte dynamisch aus windows11.partitions lesen
    local partfile="$BACKUP_DIR/windows11.partitions"

    # Quell-Größe: last-lba + 1 Sektoren
    SOURCE_LAST_LBA=$(grep '^last-lba:' "$partfile" | awk '{print $2}')
    SOURCE_SECTORS=$(( SOURCE_LAST_LBA + 1 ))
    SOURCE_SIZE_GB=$(( SOURCE_SECTORS * 512 / 1024 / 1024 / 1024 ))

    # Partitions-Layouts aus der Datei lesen
    # EFI (p1)
    P1_START=$(grep -E '^/dev/[a-z]+1 ' "$partfile" | grep -oP 'start=\s*\K[0-9]+')
    P1_SIZE=$(grep  -E '^/dev/[a-z]+1 ' "$partfile" | grep -oP 'size=\s*\K[0-9]+')
    P1_TYPE=$(grep  -E '^/dev/[a-z]+1 ' "$partfile" | grep -oP 'type=\K[^\s,]+')
    P1_UUID=$(grep  -E '^/dev/[a-z]+1 ' "$partfile" | grep -oP 'uuid=\K[^\s,]+')
    P1_NAME=$(grep  -E '^/dev/[a-z]+1 ' "$partfile" | grep -oP 'name="\K[^"]+')
    P1_ATTRS=$(grep -E '^/dev/[a-z]+1 ' "$partfile" | grep -oP 'attrs="\K[^"]+' || echo "")

    # MSR (p2)
    P2_START=$(grep -E '^/dev/[a-z]+2 ' "$partfile" | grep -oP 'start=\s*\K[0-9]+')
    P2_SIZE=$(grep  -E '^/dev/[a-z]+2 ' "$partfile" | grep -oP 'size=\s*\K[0-9]+')
    P2_TYPE=$(grep  -E '^/dev/[a-z]+2 ' "$partfile" | grep -oP 'type=\K[^\s,]+')
    P2_UUID=$(grep  -E '^/dev/[a-z]+2 ' "$partfile" | grep -oP 'uuid=\K[^\s,]+')
    P2_NAME=$(grep  -E '^/dev/[a-z]+2 ' "$partfile" | grep -oP 'name="\K[^"]+')

    # Windows (p3) - Start merken, Größe wird dynamisch berechnet
    P3_START=$(grep -E '^/dev/[a-z]+3 ' "$partfile" | grep -oP 'start=\s*\K[0-9]+')
    P3_TYPE=$(grep  -E '^/dev/[a-z]+3 ' "$partfile" | grep -oP 'type=\K[^\s,]+')
    P3_UUID=$(grep  -E '^/dev/[a-z]+3 ' "$partfile" | grep -oP 'uuid=\K[^\s,]+')
    P3_NAME=$(grep  -E '^/dev/[a-z]+3 ' "$partfile" | grep -oP 'name="\K[^"]+')

    # Recovery (p4) - Größe aus Original übernehmen
    P4_SIZE=$(grep  -E '^/dev/[a-z]+4 ' "$partfile" | grep -oP 'size=\s*\K[0-9]+')
    P4_TYPE=$(grep  -E '^/dev/[a-z]+4 ' "$partfile" | grep -oP 'type=\K[^\s,]+')
    P4_UUID=$(grep  -E '^/dev/[a-z]+4 ' "$partfile" | grep -oP 'uuid=\K[^\s,]+')
    P4_ATTRS=$(grep -E '^/dev/[a-z]+4 ' "$partfile" | grep -oP 'attrs="\K[^"]+' || echo "")

    LABEL_ID=$(grep '^label-id:' "$partfile" | awk '{print $2}')
}

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

    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 ""

    # Mindestgröße: belegter Speicher + Recovery + Puffer (~10 GB)
    local min_gb=$(( SOURCE_SIZE_GB / 4 + 10 ))
    if [[ $TARGET_SIZE_GB -lt $min_gb ]]; then
        echo -e "${RED}[FEHLER] Ziel-Device ist zu klein (< ${min_gb} GB).${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"
        SHRINK_GB=$(( TARGET_SIZE_GB - 8 ))
        echo -e "  ${YELLOW}⚠ Ziel ist kleiner → Windows wird auf ${SHRINK_GB} GB verkleinert${NC}"
    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): " answer
    if [[ "$answer" != "ja" ]]; then
        echo -e "${YELLOW}Abgebrochen.${NC}"
        exit 0
    fi
    echo ""
}

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

    local target_sectors
    target_sectors=$(lsblk -d -b -o SIZE "/dev/$TARGET_DEV" | tail -1 | tr -d ' ')
    target_sectors=$(( target_sectors / 512 ))

    # GPT: letzte 33 Sektoren reserviert
    local last_lba=$(( target_sectors - 34 ))

    # Recovery ans Ende, Größe aus Original übernehmen
    local p4_start=$(( last_lba - P4_SIZE + 1 ))
    local p3_size=$(( p4_start - P3_START ))

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

    # Attrs-Felder nur wenn vorhanden
    local p1_attrs_str=""
    [[ -n "$P1_ATTRS" ]] && p1_attrs_str=", attrs=\"$P1_ATTRS\""
    local p4_attrs_str=""
    [[ -n "$P4_ATTRS" ]] && p4_attrs_str=", attrs=\"$P4_ATTRS\""

    sfdisk "/dev/$TARGET_DEV" << EOF
label: gpt
label-id: ${LABEL_ID}
unit: sectors
first-lba: 34
sector-size: 512
/dev/${TARGET_DEV}1 : start=${P1_START}, size=${P1_SIZE}, type=${P1_TYPE}, uuid=${P1_UUID}, name="${P1_NAME}"${p1_attrs_str}
/dev/${TARGET_DEV}2 : start=${P2_START}, size=${P2_SIZE}, type=${P2_TYPE}, uuid=${P2_UUID}, name="${P2_NAME}"
/dev/${TARGET_DEV}3 : start=${P3_START}, size=${p3_size}, type=${P3_TYPE}, uuid=${P3_UUID}, name="${P3_NAME}"
/dev/${TARGET_DEV}4 : start=${p4_start}, size=${P4_SIZE}, type=${P4_TYPE}, uuid=${P4_UUID}${p4_attrs_str}
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"

    # Freien Loop-Device suchen
    local loop_dev
    loop_dev=$(losetup -f)

    trap "losetup -d '$loop_dev' 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 "$loop_dev" "$tmp_img"
    echo "y" | ntfsresize -s "${SHRINK_GB}G" "$loop_dev"

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

    losetup -d "$loop_dev"
    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}"
    ntfsfix "/dev/${TARGET_DEV}3" 2>/dev/null || true
    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
read_partition_info

echo -e "${BOLD}  Quell-Image:${NC} ${SOURCE_SIZE_GB} GB (last-lba: ${SOURCE_LAST_LBA})"
echo -e "${BOLD}  EFI-Start:${NC}   Sektor ${P1_START}, Größe ${P1_SIZE}"
echo -e "${BOLD}  Win-Start:${NC}   Sektor ${P3_START}"
echo -e "${BOLD}  Recovery:${NC}    ${P4_SIZE} Sektoren"
echo ""

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 ""