#!/bin/bash
# ==============================================================================
# karch module: disk
# Phase 1 — Partition, format, and mount target disk
# ==============================================================================

run_disk() {
    # --- Detect boot mode ----------------------------------------------------
    local boot_mode
    boot_mode="$(detect_boot_mode)"
    write_env "BOOT_MODE" "$boot_mode"

    # --- Select target disk --------------------------------------------------
    local disk_list=()
    while IFS= read -r line; do
        local dname dsize
        dname="$(echo "$line" | awk '{print $1}')"
        dsize="$(echo "$line" | awk '{print $2}')"
        disk_list+=("$dname" "${dsize}")
    done < <(lsblk -d -n -p -o NAME,SIZE,TYPE | awk '$3 == "disk" {print $1, $2}')

    if [[ ${#disk_list[@]} -eq 0 ]]; then
        dialog_msgbox "Disk" "No disks found!"
        return 1
    fi

    local target_disk
    target_disk=$(dialog_menu "Select Disk" \
        "Choose the target disk for installation.\n\nWARNING: All data on the selected disk will be DESTROYED!" \
        "${disk_list[@]}") || return 1

    write_env "DISK" "$target_disk"

    # --- Partitioning method -------------------------------------------------
    local method
    method=$(dialog_menu "Partitioning" \
        "How would you like to partition ${target_disk}?" \
        "guided" "Guided (automatic partitioning)" \
        "manual" "Manual (use cfdisk)") || return 1

    local psuf
    psuf="$(part_suffix "$target_disk")"

    if [[ "$method" == "guided" ]]; then
        _guided_partition "$target_disk" "$boot_mode" "$psuf"
    else
        _manual_partition "$target_disk" "$psuf"
    fi
}

# --- Guided partitioning ----------------------------------------------------

_guided_partition() {
    local disk="$1" boot_mode="$2" psuf="$3"

    # Confirm destructive operation
    dialog_yesno "Confirm" \
        "This will ERASE ALL DATA on ${disk}.\n\nAre you sure?" || return 1

    # Detect RAM for default swap size
    local ram_mb
    ram_mb="$(awk '/MemTotal/ {print int($2 / 1024)}' /proc/meminfo)"
    local default_swap_mb=$((ram_mb / 2))

    # --- Collect partition sizes ---------------------------------------------
    local efi_size_mb=""
    if [[ "$boot_mode" == "uefi" ]]; then
        local efi_input
        while true; do
            efi_input=$(dialog_input "EFI Partition" \
                "Size of the EFI System Partition in MiB:" \
                "1024") || return 1
            if [[ "$efi_input" =~ ^[0-9]+$ ]] && [[ "$efi_input" -ge 100 ]]; then
                break
            else
                dialog_msgbox "Invalid Input" "Please enter a numeric value (minimum 100 MiB)."
            fi
        done
        efi_size_mb="$efi_input"
    fi

    # Ask about swap
    local swap_size_mb="0"
    if dialog_yesno "Swap Partition" \
        "Create a swap partition?\n\nDetected RAM: ${ram_mb} MiB\nRecommended swap: ${default_swap_mb} MiB"; then
        local swap_input
        while true; do
            swap_input=$(dialog_input "Swap Size" \
                "Size of the swap partition in MiB:" \
                "$default_swap_mb") || return 1
            if [[ "$swap_input" =~ ^[0-9]+$ ]] && [[ "$swap_input" -gt 0 ]]; then
                break
            else
                dialog_msgbox "Invalid Input" "Please enter a positive numeric value."
            fi
        done
        swap_size_mb="$swap_input"
    fi

    # --- Create partition table ----------------------------------------------
    msg_info "Partitioning ${disk}..."

    # Wipe existing partition table
    sgdisk --zap-all "$disk" &>/dev/null
    partprobe "$disk" 2>/dev/null || true
    sleep 1

    local part_num=1

    if [[ "$boot_mode" == "uefi" ]]; then
        # EFI System Partition
        if ! sgdisk -n "${part_num}:0:+${efi_size_mb}M" -t "${part_num}:ef00" -c "${part_num}:EFI" "$disk"; then
            msg_error "Failed to create EFI partition!"
            return 1
        fi
        write_env "PART_EFI" "${disk}${psuf}${part_num}"
        part_num=$((part_num + 1))
    else
        # BIOS boot partition (1 MiB for GRUB)
        if ! sgdisk -n "${part_num}:0:+1M" -t "${part_num}:ef02" -c "${part_num}:BIOS_BOOT" "$disk"; then
            msg_error "Failed to create BIOS boot partition!"
            return 1
        fi
        write_env "PART_EFI" ""
        part_num=$((part_num + 1))
    fi

    # Swap partition (optional)
    if [[ "$swap_size_mb" -gt 0 ]]; then
        if ! sgdisk -n "${part_num}:0:+${swap_size_mb}M" -t "${part_num}:8200" -c "${part_num}:Swap" "$disk"; then
            msg_error "Failed to create swap partition!"
            return 1
        fi
        write_env "PART_SWAP" "${disk}${psuf}${part_num}"
        part_num=$((part_num + 1))
    else
        write_env "PART_SWAP" ""
    fi

    # Root partition (rest of disk)
    if ! sgdisk -n "${part_num}:0:0" -t "${part_num}:8304" -c "${part_num}:Root" "$disk"; then
        msg_error "Failed to create root partition!"
        return 1
    fi
    write_env "PART_ROOT" "${disk}${psuf}${part_num}"

    partprobe "$disk" 2>/dev/null || true
    sleep 2

    # --- Format partitions ---------------------------------------------------
    source_env
    msg_info "Formatting partitions..."

    if [[ "$boot_mode" == "uefi" && -n "${PART_EFI:-}" ]]; then
        msg_info "Formatting EFI partition: $PART_EFI"
        mkfs.fat -F 32 "$PART_EFI"
    fi

    if [[ -n "${PART_SWAP:-}" ]]; then
        msg_info "Formatting swap partition: $PART_SWAP"
        mkswap "$PART_SWAP"
    fi

    msg_info "Formatting root partition: $PART_ROOT"
    mkfs.ext4 -F "$PART_ROOT"

    # --- Mount ---------------------------------------------------------------
    _mount_partitions
}

# --- Manual partitioning ----------------------------------------------------

_manual_partition() {
    local disk="$1" psuf="$2"

    dialog_msgbox "Manual Partitioning" \
        "You will now be dropped into cfdisk to partition the disk.\n\nWhen finished, write the table and quit.\nYou will then assign partition roles."

    cfdisk "$disk"

    partprobe "$disk" 2>/dev/null || true
    sleep 2

    # --- Assign partition roles ----------------------------------------------
    source_env
    local boot_mode="${BOOT_MODE:-bios}"

    # List partitions on the selected disk
    local part_list=()
    while IFS= read -r line; do
        local pname psize
        pname="$(echo "$line" | awk '{print $1}')"
        psize="$(echo "$line" | awk '{print $2}')"
        part_list+=("$pname" "${psize}")
    done < <(lsblk -n -p -o NAME,SIZE,TYPE "$disk" | awk '$3 == "part" {print $1, $2}')

    if [[ ${#part_list[@]} -eq 0 ]]; then
        dialog_msgbox "Error" "No partitions found on ${disk}. Please partition the disk first."
        return 1
    fi

    # Assign root partition
    local root_part
    root_part=$(dialog_menu "Root Partition" \
        "Select the ROOT partition:" \
        "${part_list[@]}") || return 1
    write_env "PART_ROOT" "$root_part"

    # Ask to format root
    if dialog_yesno "Format Root" "Format ${root_part} as ext4?"; then
        mkfs.ext4 -F "$root_part"
    fi

    # Assign EFI partition (UEFI only)
    if [[ "$boot_mode" == "uefi" ]]; then
        # Filter out root partition from EFI selection
        local efi_list=()
        local i
        for (( i=0; i<${#part_list[@]}; i+=2 )); do
            [[ "${part_list[$i]}" == "$root_part" ]] && continue
            efi_list+=("${part_list[$i]}" "${part_list[$((i+1))]}")
        done

        if [[ ${#efi_list[@]} -eq 0 ]]; then
            dialog_msgbox "Error" "No remaining partitions for EFI.\nYou need at least 2 partitions (root + EFI) for UEFI."
            return 1
        fi

        local efi_part
        efi_part=$(dialog_menu "EFI Partition" \
            "Select the EFI System Partition:" \
            "${efi_list[@]}") || return 1
        write_env "PART_EFI" "$efi_part"

        if dialog_yesno "Format EFI" "Format ${efi_part} as FAT32?\n\n(Only say Yes if you created a new ESP)"; then
            mkfs.fat -F 32 "$efi_part"
        fi
    else
        write_env "PART_EFI" ""
    fi

    # Assign swap partition (optional)
    if dialog_yesno "Swap" "Do you have a swap partition?"; then
        # Filter out already-assigned partitions from swap selection
        local swap_list=()
        local j
        for (( j=0; j<${#part_list[@]}; j+=2 )); do
            [[ "${part_list[$j]}" == "$root_part" ]] && continue
            [[ "${part_list[$j]}" == "${efi_part:-}" ]] && continue
            swap_list+=("${part_list[$j]}" "${part_list[$((j+1))]}")
        done

        if [[ ${#swap_list[@]} -eq 0 ]]; then
            dialog_msgbox "Error" "No remaining partitions for swap."
            write_env "PART_SWAP" ""
        else
            local swap_part
            swap_part=$(dialog_menu "Swap Partition" \
                "Select the SWAP partition:" \
                "${swap_list[@]}") || return 1
            write_env "PART_SWAP" "$swap_part"
            mkswap "$swap_part"
        fi
    else
        write_env "PART_SWAP" ""
    fi

    # --- Mount ---------------------------------------------------------------
    source_env
    _mount_partitions
}

# --- Mount helper ------------------------------------------------------------

_mount_partitions() {
    source_env

    msg_info "Mounting filesystems..."

    mount "$PART_ROOT" /mnt

    if [[ "${BOOT_MODE:-}" == "uefi" && -n "${PART_EFI:-}" ]]; then
        mount --mkdir "$PART_EFI" /mnt/boot
    fi

    if [[ -n "${PART_SWAP:-}" ]]; then
        swapon "$PART_SWAP"
    fi

    msg_ok "Filesystems mounted."
    lsblk "$DISK"
}
