#!/bin/bash

# Skrypt do formatowania pendrive'ów poniżej 512GB na NTFS
# Pomija dysk systemowy (live USB) i obsługuje wiele urządzeń jednocześnie

set -e

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Sprawdzenie uprawnień root
if [ "$EUID" -ne 0 ]; then 
    echo -e "${RED}Musisz uruchomić skrypt jako root!${NC}"
    exit 1
fi

# Sprawdzenie czy ntfs-3g jest zainstalowany
if ! command -v mkfs.ntfs &> /dev/null; then
    echo -e "${YELLOW}Instaluję ntfs-3g...${NC}"
    if command -v dnf &> /dev/null; then
        dnf install -y ntfs-3g
    elif command -v apt &> /dev/null; then
        apt update && apt install -y ntfs-3g
    else
        echo -e "${RED}Nie można znaleźć menedżera pakietów. Zainstaluj ntfs-3g ręcznie.${NC}"
        exit 1
    fi
fi

MAX_SIZE_GB=512
MAX_SIZE_BYTES=$((MAX_SIZE_GB * 1024 * 1024 * 1024))

echo -e "${GREEN}=== Skrypt formatowania pendrive'ów ===${NC}"
echo -e "${YELLOW}UWAGA: Zostaną sformatowane WSZYSTKIE pendrive'y poniżej ${MAX_SIZE_GB}GB!${NC}"
echo ""

# Wykrycie dysku systemowego (live USB)
SYSTEM_DISK=""
MOUNT_POINT=$(findmnt -n -o SOURCE / | sed 's/[0-9]*$//')
if [ -b "$MOUNT_POINT" ]; then
    SYSTEM_DISK=$(basename "$MOUNT_POINT")
    echo -e "${GREEN}Wykryto dysk systemowy: /dev/${SYSTEM_DISK} (zostanie pominięty)${NC}"
else
    # Próba przez /proc/cmdline
    for dev in $(lsblk -ndo NAME,TYPE | grep disk | awk '{print $1}'); do
        if mount | grep -q "/dev/${dev}"; then
            SYSTEM_DISK=$dev
            echo -e "${GREEN}Wykryto dysk systemowy: /dev/${SYSTEM_DISK} (zostanie pominięty)${NC}"
            break
        fi
    done
fi

echo ""
echo -e "${YELLOW}Wykrywanie pendrive'ów...${NC}"
echo ""

# Lista urządzeń do sformatowania
DEVICES_TO_FORMAT=()

# Przeszukaj wszystkie urządzenia blokowe
for dev in $(lsblk -ndo NAME,TYPE,SIZE,TRAN | grep "disk" | awk '{print $1}'); do
    # Pomiń dysk systemowy
    if [ "$dev" = "$SYSTEM_DISK" ]; then
        echo -e "${YELLOW}Pomijam dysk systemowy: /dev/${dev}${NC}"
        continue
    fi
    
    # Sprawdź czy to USB
    TRAN=$(lsblk -ndo TRAN /dev/${dev} 2>/dev/null || echo "")
    if [ "$TRAN" != "usb" ]; then
        continue
    fi
    
    # Pobierz rozmiar w bajtach
    SIZE_BYTES=$(lsblk -ndo SIZE -b /dev/${dev})
    SIZE_GB=$((SIZE_BYTES / 1024 / 1024 / 1024))
    
    # Sprawdź czy poniżej limitu
    if [ "$SIZE_BYTES" -lt "$MAX_SIZE_BYTES" ]; then
        DEVICES_TO_FORMAT+=("/dev/${dev}")
        VENDOR=$(lsblk -ndo VENDOR /dev/${dev} 2>/dev/null | xargs || echo "Unknown")
        MODEL=$(lsblk -ndo MODEL /dev/${dev} 2>/dev/null | xargs || echo "Unknown")
        echo -e "${GREEN}✓ Znaleziono: /dev/${dev} - ${SIZE_GB}GB (${VENDOR} ${MODEL})${NC}"
    else
        echo -e "${YELLOW}✗ Pomijam /dev/${dev} - ${SIZE_GB}GB (powyżej limitu)${NC}"
    fi
done

echo ""

# Jeśli nie znaleziono urządzeń
if [ ${#DEVICES_TO_FORMAT[@]} -eq 0 ]; then
    echo -e "${RED}Nie znaleziono żadnych pendrive'ów do sformatowania!${NC}"
    exit 0
fi

# Podsumowanie
echo -e "${YELLOW}=== Podsumowanie ===${NC}"
echo -e "Urządzenia do sformatowania: ${#DEVICES_TO_FORMAT[@]}"
for dev in "${DEVICES_TO_FORMAT[@]}"; do
    SIZE_GB=$(lsblk -ndo SIZE /dev/$(basename ${dev}))
    echo "  - ${dev} (${SIZE_GB})"
done
echo ""

# Potwierdzenie
echo -e "${RED}UWAGA: Wszystkie dane na tych urządzeniach zostaną USUNIĘTE!${NC}"
read -p "Czy na pewno chcesz kontynuować? (wpisz 'TAK' żeby potwierdzić): " confirmation

if [ "$confirmation" != "TAK" ]; then
    echo -e "${YELLOW}Anulowano.${NC}"
    exit 0
fi

echo ""
echo -e "${GREEN}=== Rozpoczynam formatowanie ===${NC}"
echo ""

# Formatowanie każdego urządzenia
for dev in "${DEVICES_TO_FORMAT[@]}"; do
    DEVICE_NAME=$(basename ${dev})
    SIZE=$(lsblk -ndo SIZE ${dev})
    
    echo -e "${YELLOW}Formatuję ${dev} (${SIZE})...${NC}"
    
    # Odmontuj wszystkie partycje
    for part in $(lsblk -nlo NAME ${dev} | grep -v "^${DEVICE_NAME}$"); do
        if mountpoint -q /dev/${part} 2>/dev/null; then
            echo "  Odmontowuję /dev/${part}..."
            umount /dev/${part} 2>/dev/null || true
        fi
    done
    
    # Wyczyść istniejące sygnatury
    echo "  Czyszczę sygnatury..."
    wipefs -af ${dev} 2>/dev/null || true
    
    # Usuń wszystkie partycje
    echo "  Usuwam partycje..."
    dd if=/dev/zero of=${dev} bs=512 count=1 conv=notrunc 2>/dev/null
    
    # Poczekaj na synchronizację
    sync
    sleep 2
    
    # Utwórz nową tablicę partycji GPT
    echo "  Tworzę tablicę partycji GPT..."
    parted -s ${dev} mklabel gpt
    
    # Utwórz jedną dużą partycję NTFS (z małym marginesem na początku i końcu)
    echo "  Tworzę partycję..."
    parted -s ${dev} mkpart primary ntfs 1MiB 100%
    
    # Poczekaj na utworzenie partycji
    sleep 2
    partprobe ${dev} 2>/dev/null || true
    sleep 1
    
    # Znajdź utworzoną partycję
    PARTITION="${dev}1"
    if [ ! -b "$PARTITION" ]; then
        # Dla nvme i podobnych może być p1
        PARTITION="${dev}p1"
    fi
    
    if [ ! -b "$PARTITION" ]; then
        echo -e "${RED}  ✗ Nie można znaleźć partycji dla ${dev}${NC}"
        continue
    fi
    
    # Formatuj na NTFS z etykietą
    echo "  Formatuję na NTFS..."
    LABEL="USB_$(date +%Y%m%d)_$(echo ${DEVICE_NAME} | tail -c 3)"
    mkfs.ntfs -f -L "${LABEL}" ${PARTITION} > /dev/null 2>&1
    
    sync
    
    echo -e "${GREEN}  ✓ Sformatowano ${dev} -> NTFS (etykieta: ${LABEL})${NC}"
    echo ""
done

echo -e "${GREEN}=== Zakończono! ===${NC}"
echo -e "${YELLOW}Wszystkie pendrive'y zostały sformatowane na NTFS.${NC}"
echo ""

# Wyświetl podsumowanie
echo "Podsumowanie:"
for dev in "${DEVICES_TO_FORMAT[@]}"; do
    DEVICE_NAME=$(basename ${dev})
    SIZE=$(lsblk -ndo SIZE ${dev})
    LABEL=$(lsblk -ndo LABEL ${dev}1 2>/dev/null || lsblk -ndo LABEL ${dev}p1 2>/dev/null || echo "N/A")
    echo "  ${dev} - ${SIZE} - Etykieta: ${LABEL}"
done