#!/bin/bash

# Цвета
C_RESET='\033[0m'
C_RED='\033[0;31m'
C_GREEN='\033[0;32m'
C_YELLOW='\033[0;33m'
C_BLUE='\033[0;34m'
C_CYAN='\033[0;36m'

# Файл для сохранения состояния
CONFIG_FILE=".btrfs_setup.conf"

# --- Функции Помощники ---

info() { echo -e "${C_CYAN}INFO:${C_RESET} $1"; }
warn() { echo -e "${C_YELLOW}WARN:${C_RESET} $1"; }
error() { echo -e "${C_RED}ERROR:${C_RESET} $1"; }
success() { echo -e "${C_GREEN}SUCCESS:${C_RESET} $1"; }
cmd() { echo -e "${C_BLUE}CMD:${C_RESET} $1"; }

ask_confirm() {
    while true; do
        read -p "$(echo -e "${C_YELLOW}CONFIRM:${C_RESET} $1 (y/n): ")" yn
        case $yn in
            [Yy]* ) return 0;;
            [Nn]* ) return 1;;
            * ) echo "Ответьте 'y' или 'n'.";;
        esac
    done
}

# --- Управление Конфигурацией ---

# Загрузка переменных из файла
load_config() {
    if [ -f "$CONFIG_FILE" ]; then
        source "$CONFIG_FILE"
    fi
}

# Сохранение переменной в файл
save_config() {
    local key=$1
    local value=$2
    # Удаляем старую запись, если она есть
    sed -i "/^${key}=/d" "$CONFIG_FILE"
    # Добавляем новую
    echo "${key}='${value}'" >> "$CONFIG_FILE"
}

# --- Главы Установки ---

# Глава 0: Проверка Live-среды
step0_check_live() {
    info "--- Глава 0: Проверка Live-среды ---"
    info "Проверяю наличие 'btrfs-progs'..."
    if ! pacman -S --needed --noconfirm btrfs-progs; then
        error "Не удалось установить 'btrfs-progs'. Проверьте интернет."
        return 1
    fi
    success "Live-среда готова."
}

# Глава 1: Форматирование
step1_format() {
    info "--- Глава 1: Форматирование ---"
    read -p "EFI раздел [${DEFAULT_EFI_PART:-/dev/sda1}]: " EFI_PART
    EFI_PART=${EFI_PART:-${DEFAULT_EFI_PART:-/dev/sda1}}
    
    read -p "BTRFS раздел [${DEFAULT_BTRFS_PART:-/dev/sda2}]: " BTRFS_PART
    BTRFS_PART=${BTRFS_PART:-${DEFAULT_BTRFS_PART:-/dev/sda2}}

    if [ -z "$EFI_PART" ] || [ -z "$BTRFS_PART" ]; then
        error "Разделы не могут быть пустыми."; return 1;
    fi

    warn "!! ВНИМАНИЕ !! Это действие необратимо."
    if ask_confirm "Отформатировать $EFI_PART (fat32) и $BTRFS_PART (btrfs)?"; then
        cmd "mkfs.fat -F 32 $EFI_PART"
        mkfs.fat -F 32 "$EFI_PART" || { error "Ошибка форматирования EFI."; return 1; }
        
        cmd "mkfs.btrfs -f $BTRFS_PART"
        mkfs.btrfs -f "$BTRFS_PART" || { error "Ошибка форматирования Btrfs."; return 1; }
        
        save_config "DEFAULT_EFI_PART" "$EFI_PART"
        save_config "DEFAULT_BTRFS_PART" "$BTRFS_PART"
        success "Форматирование завершено."
    else
        info "Форматирование отменено."; return 1;
    fi
}

# Глава 2: Создание подтомов
step2_create_subvols() {
    info "--- Глава 2: Создание подтомов ---"
    if [ -z "$DEFAULT_BTRFS_PART" ]; then
        error "Сначала выполните Главу 1 (Форматирование)."; return 1;
    fi

    cmd "mount -o compress=zstd:1 $DEFAULT_BTRFS_PART /mnt"
    mount -o compress=zstd:1 "$DEFAULT_BTRFS_PART" /mnt || { error "Не удалось смонтировать /mnt."; return 1; }

    local default_vols=("@" "@home" "@srv" "@logs" "@cache" "@snapshots" "@swap")
    local SUBVOLS_RAW=${DEFAULT_SUBVOLS:-${default_vols[*]}}
    
    info "Текущий список подтомов: $SUBVOLS_RAW"
    if ask_confirm "Хотите настроить список подтомов?"; then
        read -p "Введите новый список (через пробел): " SUBVOLS_RAW
    fi

    local SUBVOLS_LIST=($SUBVOLS_RAW)
    info "Создаю подтома: ${SUBVOLS_LIST[*]}"
    for vol in "${SUBVOLS_LIST[@]}"; do
        cmd "btrfs subvolume create /mnt/$vol"
        btrfs subvolume create "/mnt/$vol"
    done
    
    save_config "DEFAULT_SUBVOLS" "$SUBVOLS_RAW"
    success "Подтома созданы."
    cmd "umount /mnt"
    umount /mnt
}

# Глава 3: Монтирование подтомов
step3_mount_all() {
    info "--- Глава 3: Монтирование подтомов ---"
    if [ -z "$DEFAULT_BTRFS_PART" ] || [ -z "$DEFAULT_SUBVOLS" ]; then
        error "Сначала выполните Главы 1 и 2."; return 1;
    fi

    local SUBVOLS_LIST=($DEFAULT_SUBVOLS)
    declare -A MOUNT_MAP
    MOUNT_MAP["@"]="/mnt"
    MOUNT_MAP["@home"]="/mnt/home"
    MOUNT_MAP["@srv"]="/mnt/srv"
    MOUNT_MAP["@logs"]="/mnt/var/log"
    MOUNT_MAP["@cache"]="/mnt/var/cache"
    MOUNT_MAP["@snapshots"]="/mnt/.snapshots"
    MOUNT_MAP["@swap"]="/mnt/swap"

    # Добавляем кастомные, если они были
    for vol in "${SUBVOLS_LIST[@]}"; do
        if [[ -z "${MOUNT_MAP[$vol]}" ]]; then
            read -p "Введите точку монтирования для '$vol' (e.g., /mnt/custom): " custom_mount
            MOUNT_MAP["$vol"]="$custom_mount"
        fi
    done

    read -p "Уровень сжатия по умолчанию [zstd:3]: " DEFAULT_COMPRESS
    DEFAULT_COMPRESS=${DEFAULT_COMPRESS:-zstd:3}
    
    info "Монтирование корня (@)..."
    local root_opts="defaults,compress=$DEFAULT_COMPRESS"
    if [[ " ${SUBVOLS_LIST[*]} " =~ " @cache " ]]; then # Корень не должен иметь noatime
        root_opts="defaults,compress=$DEFAULT_COMPRESS"
    fi
    mkdir -p "${MOUNT_MAP['@']}"
    cmd "mount -o $root_opts,subvol=@ $DEFAULT_BTRFS_PART ${MOUNT_MAP['@']}"
    mount -o "$root_opts,subvol=@" "$DEFAULT_BTRFS_PART" "${MOUNT_MAP['@']}" || { error "Ошибка монтирования корня."; return 1; }

    # Монтируем все остальное
    for vol in "${SUBVOLS_LIST[@]}"; do
        if [[ "$vol" == "@" ]]; then continue; fi
        
        local mount_point="${MOUNT_MAP[$vol]}"
        local default_opts="defaults,compress=$DEFAULT_COMPRESS"
        
        if [[ "$vol" == "@cache" || "$vol" == "@swap" ]]; then
            default_opts="defaults,noatime"
        fi
        
        read -p "Опции для '$vol' [${default_opts}]: " custom_opts
        local opts=${custom_opts:-$default_opts}
        
        info "Монтирование '$vol' в '$mount_point'..."
        mkdir -p "$mount_point"
        cmd "mount -o $opts,subvol=$vol $DEFAULT_BTRFS_PART $mount_point"
        mount -o "$opts,subvol=$vol" "$DEFAULT_BTRFS_PART" "$mount_point" || warn "Не удалось смонтировать $vol."
    done

    info "Монтирование EFI..."
    mkdir -p /mnt/efi
    cmd "mount $DEFAULT_EFI_PART /mnt/efi"
    mount "$DEFAULT_EFI_PART" /mnt/efi || { error "Ошибка монтирования EFI."; return 1; }

    success "Все подтома смонтированы."
}

# Глава 4: Swap-файл
step4_swapfile() {
    info "--- Глава 4: Swap-файл ---"
    if [[ ! " $DEFAULT_SUBVOLS " =~ " @swap " ]]; then
        warn "Подтом '@swap' не найден, пропускаю."; return 0;
    fi
    
    read -p "Размер swap-файла [4G]: " SWAP_SIZE
    SWAP_SIZE=${SWAP_SIZE:-4G}
    
    info "Создаю swap-файл /mnt/swap/swapfile..."
    cmd "btrfs filesystem mkswapfile --size $SWAP_SIZE --uuid clear /mnt/swap/swapfile"
    btrfs filesystem mkswapfile --size "$SWAP_SIZE" --uuid clear /mnt/swap/swapfile
    
    cmd "swapon /mnt/swap/swapfile"
    swapon /mnt/swap/swapfile
    save_config "SWAP_SIZE" "$SWAP_SIZE"
    success "Swap-файл создан и активирован."
}

# Глава 5: Pacstrap
step5_pacstrap() {
    info "--- Глава 5: Установка (pacstrap) ---"
    
    # Предлагаем "умный" список
    local BASE_PKGS="base linux linux-firmware nano"
    local BTRFS_UTILS=""
    if [[ " $DEFAULT_SUBVOLS " =~ " @snapshots " ]]; then
        BTRFS_UTILS="btrfs-progs grub efibootmgr snapper grub-btrfs snap-pac inotify-tools"
    else
        BTRFS_UTILS="btrfs-progs grub efibootmgr"
    fi
    
    local DEFAULT_PACKAGES="${BASE_PKGS} ${BTRFS_UTILS}"
    info "Рекомендуемый список пакетов:"
    info "$DEFAULT_PACKAGES"
    
    read -p "Введите список пакетов (Enter = принять): " PACKAGES
    PACKAGES=${PACKAGES:-$DEFAULT_PACKAGES}
    
    info "Запуск pacstrap (это займет время)..."
    cmd "pacstrap -K /mnt $PACKAGES"
    if pacstrap -K /mnt $PACKAGES; then
        save_config "INSTALLED_PACKAGES" "$PACKAGES"
        success "Базовая система установлена."
    else
        error "Ошибка pacstrap!"; return 1;
    fi
}

# Глава 6: Fstab
step6_fstab() {
    info "--- Глава 6: Fstab ---"
    info "Генерирую fstab..."
    cmd "genfstab -U /mnt >> /mnt/etc/fstab"
    genfstab -U /mnt >> /mnt/etc/fstab

    if [ -n "$SWAP_SIZE" ]; then
        info "Добавляю swap-файл в fstab..."
        cmd "echo '/swap/swapfile none swap defaults 0 0' >> /mnt/etc/fstab"
        echo "/swap/swapfile none swap defaults 0 0" >> /mnt/etc/fstab
    fi

    warn "ВАЖНО: genfstab мог 'забыть' опции сжатия."
    if ask_confirm "Открыть /mnt/etc/fstab в nano для ручной проверки?"; then
        nano /mnt/etc/fstab
    fi
    success "Fstab готов."
}

# Глава 7: Шпаргалка (Chroot)
step7_show_helper() {
    info "--- Глава 7: Шпаргалка по настройке (внутри chroot) ---"
    if [ -z "$INSTALLED_PACKAGES" ]; then
        warn "Пакеты еще не установлены (Глава 5). Шпаргалка может быть неполной."
    fi
    
    echo -e "${C_CYAN}==========================================================="
    echo "           Дальнейшие шаги (выполнять в 'arch-chroot /mnt')           "
    echo "===========================================================${C_RESET}"
    echo ""
    echo "Не забудьте про: passwd, ln -sf /usr/share/zoneinfo/..., hwclock, locale.gen, hostname, hosts"
    echo ""

    # --- Динамическая секция Snapper ---
    if [[ "$INSTALLED_PACKAGES" == *"snapper"* ]]; then
        echo -e "${C_YELLOW}### 1. Настройка Snapper ###${C_RESET}"
        warn "Рекомендация: этот шаг лучше делать ПОСЛЕ первой загрузки, а не в chroot."
        echo "cmd: snapper -c root create-config /"
        echo "cmd: rm -fdir /.snapshots"
        echo "cmd: mkdir /.snapshots"
        echo "cmd: mount -a"
        echo "cmd: nano /etc/snapper/configs/root"
        echo "     > Измени: TIMELINE_CREATE=\"no\""
        echo ""
    fi

    # --- Динамическая секция GRUB ---
    if [[ "$INSTALLED_PACKAGES" == *"grub"* ]]; then
        echo -e "${C_YELLOW}### 2. Настройка GRUB ###${C_RESET}"
        echo "cmd: grub-install --target=x86_64-efi --efi-directory=/efi --bootloader-id=Arch"
        
        if [[ "$INSTALLED_PACKAGES" == *"grub-btrfs"* ]]; then
            echo "info: Добавляем поддержку снимков в GRUB..."
            echo "cmd: nano /etc/default/grub"
            echo "     > Добавь/раскомментируй: GRUB_BTRFS_SNAPSHOT_BOOT=true"
        fi
        
        echo "cmd: grub-mkconfig -o /boot/grub/grub.cfg"
        echo ""
    fi

    # --- Динамическая секция Сервисов ---
    if [[ "$INSTALLED_PACKAGES" == *"snapper"* || "$INSTALLED_PACKAGES" == *"grub-btrfs"* ]]; then
        echo -e "${C_YELLOW}### 3. Включение сервисов ###${C_RESET}"
        if [[ "$INSTALLED_PACKAGES" == *"snapper"* ]]; then
            echo "cmd: systemctl enable snapper-timeline.timer"
            echo "cmd: systemctl enable snapper-cleanup.timer"
        fi
        if [[ "$INSTALLED_PACKAGES" == *"grub-btrfs"* ]]; then
            echo "cmd: systemctl enable grub-btrfsd.service"
        fi
        echo ""
    fi
    echo -e "${C_CYAN}===========================================================${C_RESET}"
}

# --- Главное Меню ---

main_menu() {
    while true; do
        clear
        load_config # Загружаем состояние при каждом показе меню
        
        echo -e "${C_GREEN}--- Интерактивный Btrfs-помощник (v2.0) ---${C_RESET}"
        echo "Файл состояния: $CONFIG_FILE"
        echo ""
        echo -e "  ${C_CYAN}Разделы:${C_RESET} ${DEFAULT_EFI_PART:-Не задан} | ${DEFAULT_BTRFS_PART:-Не задан}"
        echo -e "  ${C_CYAN}Подтома:${C_RESET} ${DEFAULT_SUBVOLS:-Не заданы}"
        echo -e "  ${C_CYAN}Пакеты:${C_RESET}  ${INSTALLED_PACKAGES:-Не установлены}"
        echo ""
        echo "Выберите главу:"
        echo "-----------------------------------"
        echo " 0. Проверить Live-среду (pacman)"
        echo " 1. Форматирование разделов"
        echo " 2. Создание подтомов Btrfs"
        echo " 3. Монтирование подтомов"
        echo " 4. Создание Swap-файла"
        echo "-----------------------------------"
        echo " 5. Установка системы (pacstrap)"
        echo " 6. Генерация Fstab"
        echo " 7. Показать шпаргалку (Chroot)"
        echo "-----------------------------------"
        echo " q. Выход"
        echo ""
        
        read -p "Ваш выбор: " choice
        
        # Переменная для паузы
        local pause=true
        
        case $choice in
            0) step0_check_live ;;
            1) step1_format ;;
            2) step2_create_subvols ;;
            3) step3_mount_all ;;
            4. | 4) step4_swapfile ;;
            5) step5_pacstrap ;;
            6) step6_fstab ;;
            7) step7_show_helper ;;
            q|Q) break ;;
            *) warn "Неверный выбор." ;;
        esac
        
        if [ "$pause" = true ]; then
            echo ""
            read -p "Нажмите Enter для возврата в меню..."
        fi
    done
    info "Готово. Удачи!"
}

# --- Запуск ---
main_menu