#!/bin/bash

# Rely on Omarchy for styling and error handling
source "$OMARCHY_INSTALL/helpers/all.sh"

# Ensure tzupdate is available for guessing timezone
command -v tzupdate >/dev/null || sudo pacman -Sy --noconfirm tzupdate

abort() {
  gum style "${1:-Aborted installation}"
  echo
  gum style "You can retry later by running: ./.automated_script.sh"
  exit 1
}

step() {
  clear_logo
  echo
  gum style "$1"
  echo
}

notice() {
  clear_logo
  echo
  gum spin --spinner "pulse" --title "$1" -- sleep "${2:-2}"
  echo
}

# STEP 1: KEYBOARD LAYOUT

keyboard_form() {
  step "Let's setup your machine..."
  keyboards=$'Azerbaijani|azerty
Belarusian|by
Belgian|be-latin1
Bosnian|ba
Bulgarian|bg-cp1251
Croatian|croat
Czech|cz
Danish|dk-latin1
Dutch|nl
English (UK)|uk
English (US)|us
English (US, Dvorak)|dvorak
English (US, Colemak)|colemak
Estonian|et
Finnish|fi
French|fr
French (Canada)|cf
French (Switzerland)|fr_CH
Georgian|ge
German|de
German (Switzerland)|de_CH-latin1
Greek|gr
Hebrew|il
Hungarian|hu
Icelandic|is-latin1
Irish|ie
Italian|it
Japanese|jp106
Kazakh|kazakh
Khmer (Cambodia)|khmer
Kyrgyz|kyrgyz
Lao|la-latin1
Latvian|lv
Lithuanian|lt
Macedonian|mk-utf
Norwegian|no-latin1
Polish|pl
Portuguese|pt-latin1
Portuguese (Brazil)|br-abnt2
Romanian|ro
Russian|ru
Serbian|sr-latin
Slovak|sk-qwertz
Slovenian|slovene
Spanish|es
Spanish (Latin American)|la-latin1
Swedish|sv-latin1
Tajik|tj_alt-UTF8
Turkish|trq
Ukrainian|ua'
  choice=$(printf '%s\n' "$keyboards" | cut -d'|' -f1 | gum choose --height 10 --selected "English (US)" --header "Select keyboard layout") || abort
  keyboard=$(printf '%s\n' "$keyboards" | awk -F'|' -v c="$choice" '$1==c{print $2; exit}')

  # Only attempt to load keyboard layout if we're on a real console
  # loadkeys only works on Linux virtual consoles (tty), not in terminal emulators
  if [[ $(tty 2>/dev/null) == "/dev/tty"* ]]; then
    loadkeys "$keyboard" 2>/dev/null
  fi
}

keyboard_form

# STEP 2: USER

user_form() {
  step "Let's setup your user account..."

  while true; do
    username=$(gum input --placeholder "Alphanumeric without spaces (like dhh)" --prompt.foreground="#845DF9" --prompt "Username> ") || abort

    if [[ "$username" =~ ^[a-z_][a-z0-9_-]*[$]?$ ]]; then
      break
    else
      notice "Username must be alphanumeric with no spaces" 1
    fi
  done

  while true; do
    password=$(gum input --placeholder "Used for user + root, and disk encryption when enabled" --prompt.foreground="#845DF9" --password --prompt "Password> ") || abort
    password_confirmation=$(gum input --placeholder "Must match the password you just typed" --prompt.foreground="#845DF9" --password --prompt "Confirm> ") || abort

    if [[ -n "$password" && "$password" == "$password_confirmation" ]]; then
      break
    elif [[ -z "$password" ]]; then
      notice "Your password can't be blank!" 1
    else
      notice "Passwords didn't match!" 1
    fi
  done

  # Hash the password using yescrypt
  password_hash=$(printf '%s' "$password" | openssl passwd -6 -stdin)

  full_name=$(gum input --placeholder "Used for git authentication (hit return to skip)" --prompt.foreground="#845DF9" --prompt "Full name> ")
  email_address=$(gum input --placeholder "Used for git authentication (hit return to skip)" --prompt.foreground="#845DF9" --prompt "Email address> ")

  while true; do
    hostname=$(gum input --placeholder "Letters, digits, and dashes (or return for 'omarchy')" --prompt.foreground="#845DF9" --prompt "Hostname> ")

    if [[ "$hostname" =~ ^[A-Za-z0-9]([A-Za-z0-9-]{0,61}[A-Za-z0-9])?$ ]]; then
      break
    elif [[ -z $hostname ]]; then
      hostname="omarchy"
      break
    else
      notice "Hostname must be 1-63 letters, digits, or dashes, and cannot start or end with a dash" 1
    fi
  done

  # Pick timezone
  geo_guessed_timezone=$(tzupdate -p 2>/dev/null)

  if [[ -n $geo_guessed_timezone ]]; then
    timezone=$(timedatectl list-timezones | gum choose --height 10 --selected "$geo_guessed_timezone" --header "Timezone") || abort
  else
    timezone=$(timedatectl list-timezones | gum filter --height 10 --header "Timezone") || abort
  fi
}

user_form

while true; do
  # Add manual padding since gum table -p doesn't respect padding
  echo -e "Field,Value
Username,$username
Password,$(printf "%${#password}s" | tr ' ' '*')
Full name,${full_name:-[Skipped]}
Email address,${email_address:-[Skipped]}
Hostname,$hostname
Timezone,$timezone
Keyboard,$keyboard" |
    gum table -s "," -p | sed "s/^/${PADDING_LEFT_SPACES}/"

  echo
  if gum confirm --negative "No, change it" "Does this look right?"; then
    break
  else
    keyboard_form
    user_form
  fi
done

# STEP 3: DISK

get_disk_info() {
  local device="$1"
  local size vendor model label

  # Get basic disk information
  size=$(lsblk -dno SIZE "$device" 2>/dev/null)
  vendor=$(lsblk -dno VENDOR "$device" 2>/dev/null | sed 's/ *$//')
  model=$(lsblk -dno MODEL "$device" 2>/dev/null | sed 's/ *$//')

  # Combine vendor and model, avoiding duplication
  label=""
  if [[ -n $vendor && -n $model ]]; then
    if [[ $model == *$vendor* ]]; then
      label="$model"
    else
      label="$vendor $model"
    fi
  elif [[ -n $model ]]; then
    label="$model"
  elif [[ -n $vendor ]]; then
    label="$vendor"
  fi

  # Format display string
  local display="$device"
  [[ -n $size ]] && display="$display ($size)"
  [[ -n $label ]] && display="$display - $label"

  # Append compact partition summary
  local part_summary
  part_summary=$(lsblk -nro TYPE,NAME,FSTYPE,MOUNTPOINT "$device" 2>/dev/null | \
    awk '$1=="part" { printf "%s%s%s", s, ($3==""?"unknown":$3), ($4==""?"":"("$4")"); s=", " }')
  [[ -n $part_summary ]] && display+=" [$part_summary]"

  echo "$display"
}

get_root_disk() {
  local device="$1"
  local parent

  [[ -n $device ]] || return 1

  # /run/archiso/bootmnt usually resolves to the boot partition, but the
  # installer picker works with whole disks. Walk back to the parent disk so
  # the install media itself never shows up as a wipe target.
  device=$(readlink -f "$device" 2>/dev/null || printf "%s\n" "$device")

  while true; do
    parent=$(lsblk -dno PKNAME "$device" 2>/dev/null | tail -n1)
    [[ -n $parent ]] || break
    device="/dev/$parent"
  done

  if [[ $(lsblk -dno TYPE "$device" 2>/dev/null) == "disk" ]]; then
    printf "%s\n" "$device"
  fi
}

# Compute partition device path, handling NVMe/mmcblk's pN naming.
partition_path() {
  local _disk="$1" _num="$2"
  if [[ "$_disk" == *nvme* || "$_disk" == *mmcblk* ]]; then
    echo "${_disk}p${_num}"
  else
    echo "${_disk}${_num}"
  fi
}

# BitLocker writes the "-FVE-FS-" signature at offset 3 of the partition.
detect_bitlocker() {
  local p sig
  for p in $(lsblk -lnpo NAME,TYPE "$disk" | awk '$2=="part"{print $1}'); do
    sig=$(dd if="$p" bs=1 skip=3 count=8 status=none 2>/dev/null)
    if [[ "$sig" == "-FVE-FS-" ]]; then
      echo "$p"
      return 0
    fi
  done
  return 1
}

# Look for an existing Windows ESP (FAT partition containing EFI/Microsoft) on $disk.
detect_windows_esp() {
  local p tmp_mp
  for p in $(blkid -t TYPE=vfat -o device 2>/dev/null); do
    [[ "$p" == "$disk"* ]] || continue
    tmp_mp=$(mktemp -d)
    if mount -o ro "$p" "$tmp_mp" 2>/dev/null; then
      if [[ -d "$tmp_mp/EFI/Microsoft" ]]; then
        umount "$tmp_mp"; rmdir "$tmp_mp"
        echo "$p"
        return 0
      fi
      umount "$tmp_mp"
    fi
    rmdir "$tmp_mp" 2>/dev/null || true
  done
  return 1
}

to_gb() {
  awk -v bytes="$1" 'BEGIN {printf "%.1fGB", bytes/1024/1024/1024}'
}

run_partition_install() {
  local dry=false
  [[ "${1:-}" == "dry" ]] && dry=true

  step "Checking for BitLocker on $disk"
  local bl_part
  bl_part=$(detect_bitlocker || true)
  if [[ -n "$bl_part" ]]; then
    gum style --foreground 1 "BitLocker signature detected on $bl_part."
    gum style "Suspend or disable BitLocker in Windows before installing alongside it."
    abort "Aborted: BitLocker is active on this disk."
  fi

  step "Looking for an existing Windows ESP on $disk"
  local shared_esp
  shared_esp=$(detect_windows_esp || true)

  if [[ -n "$shared_esp" ]]; then
    gum style "Found Windows ESP at $shared_esp — will install alongside."
    local tmp_mp esp_free_mib
    tmp_mp=$(mktemp -d)
    mount -o ro "$shared_esp" "$tmp_mp"
    esp_free_mib=$(df -BM --output=avail "$tmp_mp" | tail -1 | tr -dc '0-9')
    umount "$tmp_mp"; rmdir "$tmp_mp"
    if (( esp_free_mib < 100 )); then
      abort "Windows ESP $shared_esp has only ${esp_free_mib}MiB free — need at least 100MiB for Omarchy."
    fi
    gum style "Windows ESP has ${esp_free_mib}MiB free."
  else
    gum style "No Windows ESP detected on $disk — a new EFI partition will be created in free space."
  fi

  step "Analyzing free space on $disk"
  partprobe "$disk" 2>/dev/null || true
  sleep 1
  parted --script "$disk" unit B print free 2>/dev/null || lsblk -b "$disk"
  echo

  local FREE_SPACE_INFO
  FREE_SPACE_INFO=$(parted -m "$disk" unit B print free 2>/dev/null | \
    awk -F: '
      $NF ~ /free/ {
        start=$2; end=$3; size=$4;
        gsub(/B/, "", start);
        gsub(/B/, "", end);
        gsub(/B/, "", size);
        if (size+0 > max_size+0) {
          max_size=size; max_start=start; max_end=end;
        }
      }
      END {
        if (max_size+0 > 0) printf "%s %s %s\n", max_start, max_end, max_size;
      }')

  if [[ -z "$FREE_SPACE_INFO" ]]; then
    gum style --foreground 1 "No usable free space detected on $disk."
    gum style "Partition install uses unallocated free space, not a pre-created partition."
    gum style "Open the partition tool, delete/shrink partitions until the target area says Free space, then leave it unallocated."
    gum confirm --affirmative "Back" --negative "Back" "Return to installation mode?" >/dev/null 2>&1 || true
    return 1
  fi

  local FREE_START_B FREE_END_B FREE_SIZE_B
  FREE_START_B=$(echo "$FREE_SPACE_INFO" | awk '{print $1}')
  FREE_END_B=$(echo "$FREE_SPACE_INFO"   | awk '{print $2}')
  FREE_SIZE_B=$(echo "$FREE_SPACE_INFO"  | awk '{print $3}')

  local mib=$((1024 * 1024)) gib=$((1024 * 1024 * 1024))
  local ALIGN=$mib
  local EFI_SIZE_B=$((2 * gib))
  local MIN_ROOT_B=$((32 * gib))    # 32 GiB minimum for the root filesystem

  local EFI_START_B=0 EFI_END_B=0 ROOT_START_B ROOT_END_B
  if [[ -z "$shared_esp" ]]; then
    EFI_START_B=$(( (FREE_START_B + ALIGN - 1) / ALIGN * ALIGN ))
    EFI_END_B=$((EFI_START_B + EFI_SIZE_B))
    ROOT_START_B=$(( (EFI_END_B + 1 + ALIGN - 1) / ALIGN * ALIGN ))
  else
    ROOT_START_B=$(( (FREE_START_B + ALIGN - 1) / ALIGN * ALIGN ))
  fi
  ROOT_END_B=$((FREE_END_B / ALIGN * ALIGN))

  if (( ROOT_END_B <= ROOT_START_B )); then
    gum style --foreground 1 "Not enough aligned free space for partitions after alignment."
    gum confirm --affirmative "Back" --negative "Back" "Return to installation mode?" >/dev/null 2>&1 || true
    return 1
  fi

  local ROOT_MAX_B=$((ROOT_END_B - ROOT_START_B))
  if (( ROOT_MAX_B < MIN_ROOT_B )); then
    local required_b=$MIN_ROOT_B
    [[ -z "$shared_esp" ]] && required_b=$((required_b + EFI_SIZE_B))
    gum style --foreground 1 "Usable root space ($(to_gb $ROOT_MAX_B)) is below the 32GB minimum for Omarchy."
    gum style "This disk needs at least $(to_gb $required_b) unallocated in this free region."
    gum style "Open the partition tool, delete/shrink partitions until the target area says Free space, then leave it unallocated."
    gum confirm --affirmative "Back" --negative "Back" "Return to installation mode?" >/dev/null 2>&1 || true
    return 1
  fi

  local ROOT_MAX_GIB=$((ROOT_MAX_B / gib))
  local ROOT_SIZE_GIB
  while true; do
    step "Choose Omarchy size"
    gum style "Free region: $(to_gb $FREE_START_B) — $(to_gb $FREE_END_B)  ($(to_gb $FREE_SIZE_B))"
    if [[ -z "$shared_esp" ]]; then
      gum style "2GiB will be used for a new EFI partition."
    else
      gum style "Using existing ESP: $shared_esp"
    fi
    gum style "Root filesystem can be 32–${ROOT_MAX_GIB}GiB. Leave the default to use all available space."
    echo

    ROOT_SIZE_GIB=$(gum input \
      --value "$ROOT_MAX_GIB" \
      --placeholder "32-${ROOT_MAX_GIB}" \
      --prompt.foreground="#845DF9" \
      --prompt "Omarchy size in GiB> ") || return 1

    if [[ $ROOT_SIZE_GIB =~ ^[0-9]+$ ]] && (( ROOT_SIZE_GIB >= 32 && ROOT_SIZE_GIB <= ROOT_MAX_GIB )); then
      break
    fi

    notice "Enter a whole number from 32 to ${ROOT_MAX_GIB}" 1
  done

  ROOT_END_B=$((ROOT_START_B + ROOT_SIZE_GIB * gib))

  step "Proposed layout on $disk"
  gum style "Free region: $(to_gb $FREE_START_B) — $(to_gb $FREE_END_B)  ($(to_gb $FREE_SIZE_B))"
  if [[ -z "$shared_esp" ]]; then
    gum style "  New EFI partition:  $(to_gb $EFI_START_B) — $(to_gb $EFI_END_B) (2GB FAT32)"
  else
    gum style "  Shared ESP:        $shared_esp (Windows, mounted at /efi)"
  fi
  gum style "  New root (Btrfs):  $(to_gb $ROOT_START_B) — $(to_gb $ROOT_END_B) ($(to_gb $((ROOT_END_B - ROOT_START_B))))"
  echo

  gum confirm --affirmative "Yes, partition" --negative "Cancel" \
    "Create Omarchy partitions in the free space on $disk?" || return 1

  echo
  if gum confirm \
    --affirmative "Yes, encrypt" \
    --negative "No, leave unencrypted" \
    "Encrypt the Omarchy root partition with LUKS2?"; then
    encrypt_installation=true
  else
    encrypt_installation=false
  fi

  if [[ "$encrypt_installation" == "true" && -n "$shared_esp" ]]; then
    abort "Encrypted install on a shared Windows ESP is not yet supported. Please choose unencrypted, or use Full Install."
  fi

  local esp_mount_in_target
  if [[ "$encrypt_installation" == "true" ]]; then
    esp_mount_in_target=/boot
  else
    esp_mount_in_target=/efi
  fi

  # Detect T2 Mac and set appropriate kernel
  if lspci -nn 2>/dev/null | grep -q "106b:180[12]"; then
    kernel_choice="linux-t2"
  else
    kernel_choice="linux"
  fi

  local last_part efi_part_num root_part_num
  last_part=$(lsblk -n -o PARTN "$disk" 2>/dev/null | grep -E '^[0-9]+$' | sort -n | tail -1)
  if [[ -z "$shared_esp" ]]; then
    efi_part_num=$(( ${last_part:-0} + 1 ))
    root_part_num=$(( efi_part_num + 1 ))
  else
    efi_part_num=0   # not created by us
    root_part_num=$(( ${last_part:-0} + 1 ))
  fi

  local efi_dev root_partition_device
  if [[ -n "$shared_esp" ]]; then
    efi_dev="$shared_esp"
  else
    efi_dev=$(partition_path "$disk" "$efi_part_num")
  fi
  root_partition_device=$(partition_path "$disk" "$root_part_num")

  if $dry; then
    gum style "[dry] would create root partition at $root_partition_device"
    [[ -z "$shared_esp" ]] && gum style "[dry] would create EFI partition at $efi_dev"
    gum style "[dry] encrypted=$encrypt_installation kernel=$kernel_choice"
  else
    step "Creating partitions on $disk"
    if [[ -z "$shared_esp" ]]; then
      parted --script "$disk" mkpart primary fat32 "${EFI_START_B}B" "${EFI_END_B}B"
      parted --script "$disk" set "$efi_part_num" esp on
      parted --script "$disk" name "$efi_part_num" "OMARCHY_EFI"
    fi
    parted --script "$disk" mkpart primary btrfs "${ROOT_START_B}B" "${ROOT_END_B}B"
    parted --script "$disk" name "$root_part_num" "OMARCHY_ROOT"
    partprobe "$disk"
    sync
    sleep 2

    local i
    for i in 1 2 3 4 5; do
      if [[ -b "$root_partition_device" ]] && { [[ -n "$shared_esp" ]] || [[ -b "$efi_dev" ]]; }; then
        break
      fi
      sleep 1
    done

    local root_mapper
    if [[ "$encrypt_installation" == "true" ]]; then
      step "Setting up LUKS2 on $root_partition_device"
      printf "%s" "$password" | cryptsetup luksFormat --type luks2 --batch-mode "$root_partition_device" -
      printf "%s" "$password" | cryptsetup open "$root_partition_device" omarchy_root -
      root_mapper="/dev/mapper/omarchy_root"
    else
      wipefs -af "$root_partition_device"
      root_mapper="$root_partition_device"
    fi

    step "Creating Btrfs filesystem and subvolumes"
    mkfs.btrfs -f -L OMARCHY "$root_mapper"
    for i in 1 2 3 4 5; do
      [[ -b "$root_mapper" ]] && break
      sleep 1
    done

    mkdir -p /mnt/btrfs-root
    mount "$root_mapper" /mnt/btrfs-root
    btrfs subvolume create /mnt/btrfs-root/@
    btrfs subvolume create /mnt/btrfs-root/@home
    btrfs subvolume create /mnt/btrfs-root/@log
    btrfs subvolume create /mnt/btrfs-root/@pkg
    umount /mnt/btrfs-root
    rmdir /mnt/btrfs-root

    mount -o noatime,compress=zstd,subvol=@ "$root_mapper" /mnt
    mkdir -p /mnt/home /mnt/var/log /mnt/var/cache/pacman/pkg /mnt"$esp_mount_in_target"
    mount -o noatime,compress=zstd,subvol=@home "$root_mapper" /mnt/home
    mount -o noatime,compress=zstd,subvol=@log  "$root_mapper" /mnt/var/log
    mount -o noatime,compress=zstd,subvol=@pkg  "$root_mapper" /mnt/var/cache/pacman/pkg

    if [[ -z "$shared_esp" ]]; then
      mkfs.fat -F32 -n OMARCHY_EFI "$efi_dev"
    fi
    mount "$efi_dev" /mnt"$esp_mount_in_target"
  fi

  clear

  echo "$full_name" >user_full_name.txt
  echo "$email_address" >user_email_address.txt
  printf '%s\n' "$encrypt_installation" >user_encrypt_installation.txt

  local password_escaped password_hash_escaped username_escaped
  password_escaped=$(echo -n "$password" | jq -Rsa)
  password_hash_escaped=$(echo -n "$password_hash" | jq -Rsa)
  username_escaped=$(echo -n "$username" | jq -Rsa)

  local credentials_encryption_line=""
  if [[ "$encrypt_installation" == "true" ]]; then
    credentials_encryption_line="    \"encryption_password\": $password_escaped,"
  fi

  cat <<-_EOF_ >user_credentials.json
{
$credentials_encryption_line
    "root_enc_password": $password_hash_escaped,
    "users": [
        {
            "enc_password": $password_hash_escaped,
            "groups": [],
            "sudo": true,
            "username": $username_escaped
        }
    ]
}
_EOF_

  cat <<-_EOF_ >user_configuration.json
{
    "app_config": null,
    "archinstall-language": "English",
    "auth_config": {},
    "audio_config": { "audio": "pipewire" },
    "bootloader_config": { "bootloader": "Limine", "uki": false, "removable": false },
    "custom_commands": [],
    "disk_config": {
        "config_type": "pre_mounted_config",
        "mountpoint": "/mnt"
    },
    "hostname": "$hostname",
    "kernels": [ "$kernel_choice" ],
    "network_config": { "type": "iso" },
    "ntp": true,
    "parallel_downloads": 8,
    "script": null,
    "services": [],
    "swap": true,
    "timezone": "$timezone",
    "locale_config": {
        "kb_layout": "$keyboard",
        "sys_enc": "UTF-8",
        "sys_lang": "en_US.UTF-8"
    },
    "mirror_config": {
        "custom_repositories": [],
        "custom_servers": [
            {"url": "https://mirror.omarchy.org/\$repo/os/\$arch"},
            {"url": "https://mirror.rackspace.com/archlinux/\$repo/os/\$arch"},
            {"url": "https://geo.mirror.pkgbuild.com/\$repo/os/\$arch"}
        ],
        "mirror_regions": {},
        "optional_repositories": []
    },
    "packages": [
        "base-devel",
        "git",
        "omarchy-keyring",
        "omarchy",
        "omarchy-installer"
    ],
    "profile_config": {
        "gfx_driver": null,
        "greeter": null,
        "profile": {}
    },
    "version": "3.0.9"
}
_EOF_

  local luks_uuid_json="null" luks_uuid
  if [[ "$encrypt_installation" == "true" ]]; then
    if $dry; then
      luks_uuid="00000000-0000-0000-0000-000000000000"
    else
      luks_uuid=$(blkid -s UUID -o value "$root_partition_device")
    fi
    luks_uuid_json="\"$luks_uuid\""
  fi

  cat >/root/protected_install.json <<JSON
{
  "esp_mount": "$esp_mount_in_target",
  "esp_path": "/EFI/Omarchy",
  "luks_uuid": $luks_uuid_json,
  "root_device": "$root_partition_device",
  "kernel": "$kernel_choice"
}
JSON

  if $dry; then
    echo -e "User Configuration:"
    cat user_configuration.json
    echo -e "\n\nUser Credentials:"
    cat user_credentials.json
    echo -e "\n\nUser Full Name:"
    cat user_full_name.txt
    echo -e "\nUser Email Address:"
    cat user_email_address.txt
    echo -e "\nDisk Encryption:"
    cat user_encrypt_installation.txt
    echo -e "\nProtected Install Intent:"
    cat /root/protected_install.json
    rm -f user_configuration.json user_credentials.json user_full_name.txt user_email_address.txt user_encrypt_installation.txt /root/protected_install.json
  fi
}

disk_form() {
  step "Let's select where to install Omarchy..."

  # Don't offer the install media as an option (Arch ISO mounts it here)
  local boot_source
  local exclude_disk
  boot_source=$(findmnt -no SOURCE /run/archiso/bootmnt 2>/dev/null || true)
  exclude_disk=$(get_root_disk "$boot_source")

  # List all installable disks, excluding the boot device if present
  local available_disks
  available_disks=$(
    lsblk -dpno NAME,TYPE |
      awk '$2=="disk"{print $1}' |
      grep -E '/dev/(sd|hd|vd|nvme|mmcblk|xv)' |
      { if [[ -n "$exclude_disk" ]]; then grep -Fvx "$exclude_disk"; else cat; fi; }
  )

  # Get available disks and format them with info
  local disk_options=""
  while IFS= read -r device; do
    if [[ -n "$device" ]]; then
      disk_info=$(get_disk_info "$device")
      disk_options="$disk_options$disk_info"$'\n'
    fi
  done <<<"$available_disks"

  selected_display=$(echo "$disk_options" | gum choose --header "Select install disk") || abort
  disk=$(echo "$selected_display" | awk '{print $1}')
}

disk_form

# STEP 4: INSTALL MODE

open_partition_tool() {
  step "Partition tool for $disk"
  gum style "Create unallocated free space for Omarchy, then write changes and quit."
  gum style --foreground 8 "Do not create an Omarchy partition here — leave the target area as Free space."
  gum style --foreground 8 "Tip: free-space install needs at least 32GB unallocated, plus 2GB more if no ESP already exists."
  echo
  gum confirm --affirmative "Open cfdisk" --negative "Back" "Open partition tool for $disk?" || return 0

  clear
  cfdisk "$disk" || true
  partprobe "$disk" 2>/dev/null || true
  sync
  sleep 1
}

install_mode_form() {
  step "Installation mode"
  install_mode=$(gum choose --header "How should Omarchy be installed on $disk?" \
    "Full disk install" \
    "Free space install (alongside existing data)" \
    "Open partition tool") || abort
}

while true; do
  install_mode_form

  case "$install_mode" in
    "Full disk install")
      break
      ;;
    "Free space install (alongside existing data)")
      if run_partition_install "$@"; then
        exit 0
      fi
      ;;
    "Open partition tool")
      open_partition_tool
      ;;
  esac
done

confirm_disk_overwrite() {
  local mode="encrypted"
  local affirmative confirm_status

  while true; do
    clear_logo
    echo
    gum style "Everything will be overwritten. There is no recovery possible."

    if [[ $mode == "encrypted" ]]; then
      gum style --foreground 8 "Press Ctrl+C for unencrypted install."
      affirmative="Yes, install"
    else
      affirmative="Yes, install without encryption"
    fi

    echo
    gum confirm --affirmative "$affirmative" --negative "No, change it" "Confirm overwriting ${disk}"
    confirm_status=$?

    case $confirm_status in
      0)
        [[ $mode == "encrypted" ]] && encrypt_installation=true || encrypt_installation=false
        return 0
        ;;
      1)
        return 1
        ;;
      130)
        [[ $mode == "encrypted" ]] && mode="unencrypted" || mode="encrypted"
        ;;
      *)
        abort
        ;;
    esac
  done
}

while ! confirm_disk_overwrite; do
  disk_form
done

clear

# Save user full name and email address, so the autostart can
# set them as OMARCHY_USER_NAME and OMARCHY_USER_EMAIL
echo "$full_name" >user_full_name.txt
echo "$email_address" >user_email_address.txt

# Escape username and password for safe inclusion
password_escaped=$(echo -n "$password" | jq -Rsa)
password_hash_escaped=$(echo -n "$password_hash" | jq -Rsa)
username_escaped=$(echo -n "$username" | jq -Rsa)

if [[ $encrypt_installation == true ]]; then
  credentials_encryption_line="    \"encryption_password\": $password_escaped,"
else
  credentials_encryption_line=""
fi

cat <<-_EOF_ >user_credentials.json
{
$credentials_encryption_line
    "root_enc_password": $password_hash_escaped,
    "users": [
        {
            "enc_password": $password_hash_escaped,
            "groups": [],
            "sudo": true,
            "username": $username_escaped
        }
    ]
}
_EOF_

printf '%s\n' "$encrypt_installation" >user_encrypt_installation.txt

# Setup partition layout
disk_size=$(lsblk -bdno SIZE "$disk")
mib=$((1024 * 1024))
gib=$((mib * 1024))
disk_size_in_mib=$((disk_size / mib * mib)) # Rounds to nearest MiB

gpt_backup_reserve=$((mib))
boot_partition_start=$((mib))
boot_partition_size=$((2 * gib))

main_partition_start=$((boot_partition_size + boot_partition_start))
main_partition_size=$((disk_size_in_mib - main_partition_start - gpt_backup_reserve))

# Detect T2 Mac and set appropriate kernel
if lspci -nn 2>/dev/null | grep -q "106b:180[12]"; then
  kernel_choice="linux-t2"
else
  kernel_choice="linux"
fi

if [[ $encrypt_installation == true ]]; then
  disk_encryption_config=$(cat <<-_EOF_
,
        "disk_encryption": {
            "encryption_type": "luks",
            "lvm_volumes": [],
            "iter_time": 2000,
            "partitions": [ "8c2c2b92-1070-455d-b76a-56263bab24aa" ],
            "encryption_password": $password_escaped
        }
_EOF_
)
else
  disk_encryption_config=""
fi

cat <<-_EOF_ >user_configuration.json
{
    "app_config": null,
    "archinstall-language": "English",
    "auth_config": {},
    "audio_config": { "audio": "pipewire" },
    "bootloader_config": { "bootloader": "Limine", "uki": false, "removable": false },
    "custom_commands": [],
    "disk_config": {
        "btrfs_options": { 
          "snapshot_config": {
            "type": "Snapper"
          }
        },
        "config_type": "default_layout",
        "device_modifications": [
            {
                "device": "$disk",
                "partitions": [
                    {
                        "btrfs": [],
                        "dev_path": null,
                        "flags": [ "boot", "esp" ],
                        "fs_type": "fat32",
                        "mount_options": [],
                        "mountpoint": "/boot",
                        "obj_id": "ea21d3f2-82bb-49cc-ab5d-6f81ae94e18d",
                        "size": {
                            "sector_size": { "unit": "B", "value": 512 },
                            "unit": "B",
                            "value": $boot_partition_size
                        },
                        "start": {
                            "sector_size": { "unit": "B", "value": 512 },
                            "unit": "B",
                            "value": $boot_partition_start
                        },
                        "status": "create",
                        "type": "primary"
                    },
                    {
                        "btrfs": [
                            { "mountpoint": "/", "name": "@" },
                            { "mountpoint": "/home", "name": "@home" },
                            { "mountpoint": "/var/log", "name": "@log" },
                            { "mountpoint": "/var/cache/pacman/pkg", "name": "@pkg" }
                        ],
                        "dev_path": null,
                        "flags": [],
                        "fs_type": "btrfs",
                        "mount_options": [ "compress=zstd" ],
                        "mountpoint": null,
                        "obj_id": "8c2c2b92-1070-455d-b76a-56263bab24aa",
                        "size": {
                            "sector_size": { "unit": "B", "value": 512 },
                            "unit": "B",
                            "value": $main_partition_size
                        },
                        "start": {
                            "sector_size": { "unit": "B", "value": 512 },
                            "unit": "B",
                            "value": $main_partition_start
                        },
                        "status": "create",
                        "type": "primary"
                    }
                ],
                "wipe": true
            }
        ]$disk_encryption_config
    },
    "hostname": "$hostname",
    "kernels": [ "$kernel_choice" ],
    "network_config": { "type": "iso" },
    "ntp": true,
    "parallel_downloads": 8,
    "script": null,
    "services": [],
    "swap": true,
    "timezone": "$timezone",
    "locale_config": {
        "kb_layout": "$keyboard",
        "sys_enc": "UTF-8",
        "sys_lang": "en_US.UTF-8"
    },
    "mirror_config": {
        "custom_repositories": [],
        "custom_servers": [
            {"url": "https://mirror.omarchy.org/\$repo/os/\$arch"},
            {"url": "https://mirror.rackspace.com/archlinux/\$repo/os/\$arch"},
            {"url": "https://geo.mirror.pkgbuild.com/\$repo/os/\$arch"}
        ],
        "mirror_regions": {},
        "optional_repositories": []
    },
    "packages": [
        "base-devel",
        "git",
        "omarchy-keyring",
        "omarchy",
        "omarchy-installer"
    ],
    "profile_config": {
        "gfx_driver": null,
        "greeter": null,
        "profile": {}
    },
    "version": "3.0.9"
}
_EOF_

if [[ $1 == "dry" ]]; then
  echo -e "User Configuration:"
  cat user_configuration.json

  echo -e "\n\nUser Credentials:"
  cat user_credentials.json

  echo -e "\n\nUser Full Name:"
  cat user_full_name.txt

  echo -e "\nUser Email Address:"
  cat user_email_address.txt

  echo -e "\nDisk Encryption:"
  cat user_encrypt_installation.txt

  rm user_configuration.json user_credentials.json user_full_name.txt user_email_address.txt user_encrypt_installation.txt
fi